59 lines
1.3 KiB
Python
59 lines
1.3 KiB
Python
|
# Example of a prefix tree implementation in python
|
||
|
|
# Trie data-structure implementation
|
||
|
|
# https://en.wikipedia.org/wiki/Trie
|
||
|
|
# original date: 2021-12-03
|
||
|
|
|
||
|
|
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 build_trie(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 letter in nodes:
|
||
|
|
nodes[letter] = build_trie(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_in_trie(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 = build_trie(words)
|
||
|
|
# test matching prefix "ma"
|
||
|
|
sugg = search_in_trie(res, 'ap')
|
||
|
|
print(sugg)
|
||
|
|
|
||
|
|
main()
|