python trie

# A very simple trie to put together quickly for coding problems
WORD_KEY = '$'

def build_trie(words, trie = dict()):
  for word in words: # Add each word - letter by letter
    node = trie
    for letter in word: 
      node = node.setdefault(letter, {}) # node = node[letter].child
    node[WORD_KEY] = word # Mark "leaf" with word
  return trie

def exists(trie, word, index = 0):
  if index >= len(word): return trie.get(WORD_KEY, None) == word
  return trie.get(WORD_KEY, None) == word or (word[index] in trie and exists(trie[word[index]], word, index + 1))



# API
trie = build_trie(["hello", "hello again", "bye"])
assert exists(trie, "hello")
assert exists(trie, "hello again")
assert exists(trie, "bye")
assert not exists(trie, "")
assert not exists(trie, "hel")
trie = build_trie(["bye again", "will miss you"], trie)
assert exists(trie, "hello")
assert exists(trie, "hello again")
assert exists(trie, "bye")
assert exists(trie, "bye again")
assert not exists(trie, "")
assert not exists(trie, "hel")
assert not exists(trie, "hello ")
assert not exists(trie, "hello ag")
assert not exists(trie, "byebye")
Muddy Moose