Trie (Prefix Tree) Implementation in Python

trie.py

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())