Skip to content

Instantly share code, notes, and snippets.

@osori osori/trie.py
Last active Oct 26, 2019

Embed
What would you like to do?
Trie implementation in Python
class Node(object):
"""
A single node of a trie.
Children of nodes are defined in a dictionary
where each (key, value) pair is in the form of
(Node.key, Node() object).
"""
def __init__(self, key, data=None):
self.key = key
self.data = data # data is set to None if node is not the final char of string
self.children = {}
class Trie(object):
"""
A simple Trie with insert, search, and starts_with methods.
"""
def __init__(self):
self.head = Node(None)
"""
Inserts string in the trie.
"""
def insert(self, string):
curr_node = self.head
for char in string:
if char not in curr_node.children:
curr_node.children[char] = Node(char)
curr_node = curr_node.children[char]
# When we have reached the end of the string, set the curr_node's data to string.
# This also denotes that curr_node represents the final character of string.
curr_node.data = string
"""
Returns if string exists in the trie
"""
def search(self, string):
curr_node = self.head
for char in string:
if char in curr_node.children:
curr_node = curr_node.children[char]
else:
return False
# Reached the end of string,
# If curr_node has data (i.e. curr_node is not None), string exists in the trie
if (curr_node.data != None):
return True
"""
Returns a list of words in the trie
that starts with the given prefix.
"""
def starts_with(self, prefix):
curr_node = self.head
result = []
subtrie = None
# Locate the prefix in the trie,
# and make subtrie point to prefix's last character Node
for char in prefix:
if char in curr_node.children:
curr_node = curr_node.children[char]
subtrie = curr_node
else:
return None
# Using BFS, traverse through the prefix subtrie,
# and look for nodes with non-null data fields.
queue = list(subtrie.children.values())
while queue:
curr = queue.pop()
if curr.data != None:
result.append(curr.data)
queue += list(curr.children.values())
return result
# Test
t = Trie()
words = ["romane", "romanus", "romulus", "ruben", 'rubens', 'ruber', 'rubicon', 'ruler']
for word in words:
t.insert(word)
print(t.search("romulus"))
print(t.search("ruler"))
print(t.search("rulere"))
print(t.search("romunus"))
print(t.starts_with("ro"))
print(t.starts_with("rube"))
@HongsupOH

This comment has been minimized.

Copy link

HongsupOH commented Oct 26, 2019

안녕하세요 글 잘읽었습니다. 다름아니라 starts_with에서 queue를 쓰셨는데 curr = queue.pop()으로 하셨네요. queue.pop(0)으로 바꾸시는게 맞은것 같아서요. 물론 답은 같은 답이 나오기는 합니다. 어쨌든 잘읽었습니다.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.