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