Last active
December 10, 2023 12:13
-
-
Save osori/d0200b9bb7665d6a69da61b431e4077f to your computer and use it in GitHub Desktop.
Trie implementation in Python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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")) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
안녕하세요 글 잘읽었습니다. 다름아니라 starts_with에서 queue를 쓰셨는데 curr = queue.pop()으로 하셨네요. queue.pop(0)으로 바꾸시는게 맞은것 같아서요. 물론 답은 같은 답이 나오기는 합니다. 어쨌든 잘읽었습니다.