Example of a prefix tree implementation in python

0_prefix_tree.py

import fileinput

def load_from_stdin():
  inp = fileinput.input()
  words = []
  for line in inp:
    word = line.strip()
    # non case sensitive
    words.append(word.lower())
  return words

# create the trie data structure
def make(words):
  if words == []:
    return {}
  nodes = {}
  for w in words:
    if w[0] not in nodes:
      nodes[w[0]] = []
    if len(w[1:]) > 0:
      nodes[w[0]].append(w[1:])
  for i,letter in enumerate(nodes):
    nodes[letter] = make(nodes[letter])
  return nodes

def get_sub_tree(nodes, pattern):
  if pattern == '': return nodes
  if pattern[0] not in nodes: return {}
  return get_sub_tree(nodes[pattern[0]], pattern[1:])

def get_strings(nodes):
  words = []
  for k in nodes.keys():
    if nodes[k] == {}:
      words += [ k ]
    else:
      res = get_strings(nodes[k])
      words += [ k + x for x in res]
  return words

def search(nodes, pattern):
  nodes = get_sub_tree(nodes, pattern)
  words = get_strings(nodes)
  return [ pattern + x for x in words ]

def main():
  words = load_from_stdin()
  res = make(words)
  # test matching prefix "ma"
  sugg = search(res, 'ma')
  print(sugg)

main()

fruits.txt

apple
apricot
avocado
banana
berry
blackberry
blood orange
blueberry
boysenberry
breadfruit
cantaloupe
cherry
citron
citrus
coconut
crabapple
cranberry
current
date
dragon fruit
durian
elderberry
fig
grape
grapefruit
guava
honeydew
jackfruit
kiwi
kumquat
lemon
lime
lingonberry
loquat
lychee
mandarin orange
mango
marionberry
melon
mulberry
nectarine
orange
papaya
passion fruit
peach
pear
persimmon
pineapple
plantain
plum
pluot
pomegranate
pomelo
prune
quince
raisin
raspberry
star fruit
strawberry
tangelo
tangerine
ugli fruit
watermelon