-
Notifications
You must be signed in to change notification settings - Fork 0
/
trie.py
87 lines (71 loc) · 2.69 KB
/
trie.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
"""
CS 5008
Fall 2023
Final Research Project
Marian Padron
Trie file, contains Node class to represent nodes in trie data structure and Trie class for creating the data
structure.
"""
class Node:
"""Node class, for creating nodes within trie data structure."""
def __init__(self):
"""Initialize with key, value, children, and number of children."""
self.key = None
self.value = None
self.children = {}
self.num_children = 0
class Trie:
""" Trie class for representing a trie data structure."""
def __init__(self):
"""Initialize with an empty root node."""
self.root = Node()
def insert(self, word, value):
"""
Insertion method for trie object, retuns -1 if not found.
:param word: string key
:param value: value for the key
:return: None
"""
# Begin at top of the tree
current_word = word
current_node = self.root
while len(current_word) > 0:
# If letter in node's children set current node to found node
if current_word[0] in current_node.children:
current_node = current_node.children[current_word[0]]
current_word = current_word[1:] # keep traversing the letters in the word
# If letter doesn't exist in children create a new node
else:
new_node = Node()
new_node.key = current_word[0]
# If at the end of the word assign value
if len(current_word) == 1:
new_node.value = value
# Update values
current_node.children[current_word[0]] = new_node
current_node.num_children += 1
current_node = new_node
current_word = current_word[1:]
# if word already exists, update its value
current_node.value = value
def lookup(self, word):
"""
Search method for trie object.
:param word: string key to find
:return: the value at the key
"""
current_word = word
current_node = self.root
# Traverse nodes letter by letter in word
while len(current_word) > 0:
if current_word[0] in current_node.children:
current_node = current_node.children[current_word[0]]
current_word = current_word[1:]
# If no children nodes with letter or no value at node
else:
print("not in trie" + current_word)
return -1
if current_node.value is None:
print("no value in this node")
return -1
return current_node.value