Skip to content

Instantly share code, notes, and snippets.

@theabbie
Created September 20, 2022 05:46
Show Gist options
  • Save theabbie/211a00b4bc72189102a3ad1dc8dbf616 to your computer and use it in GitHub Desktop.
Save theabbie/211a00b4bc72189102a3ad1dc8dbf616 to your computer and use it in GitHub Desktop.
Trie (Prefix Tree) Implementation in Python
class Node:
def __init__(self):
self.child = {}
self.end = False
class Trie:
def __init__(self):
self.root = Node()
def insert(self, s):
curr = self.root
for c in s:
if c not in curr.child:
curr.child[c] = Node()
curr = curr.child[c]
curr.end = True
def getAll(self, node = None):
node = node or self.root
res = []
if node.end:
res.append("")
for c in node.child:
res.extend([c + s for s in self.getAll(node.child[c])])
return res
def check(self, s):
curr = self.root
for c in s:
if c not in curr.child:
return None
curr = curr.child[c]
return curr
def startswith(self, s):
res = self.check(s)
if not res:
return []
return [s + suffix for suffix in self.getAll(res)]
def exists(self, s):
res = self.check(s)
return res and res.end
trie = Trie()
trie.insert("cat")
trie.insert("catch")
trie.insert("pen")
trie.insert("pencil")
print(trie.exists("pen"))
print(trie.exists("pe"))
print(trie.startswith("cat"))
print(trie.getAll())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment