python_workspace/ternary_tree.py at main · johnluo92/python_workspace · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
#!/usr/local/bin/python3

"""
A ternary search tree implementation, inspired by:
http://www.drdobbs.com/database/ternary-search-trees/184410528 and
https://lukaszwrobel.pl/blog/ternary-search-tree/
https://github.com/djtrack16/tst/blob/master/ternarysearchtree.py
"""

class Node:
    lo = None
    hi = None
    eq = None
    endpoint = False
    def __init__(self, char):
        self.char = char
    def __repr__(self):
        # useful in debugging
        return ''.join(['[', self.char,
                        ('' if not self.endpoint else ' <end>'),
                        ('' if self.lo is None else ' lo: ' + self.lo.__repr__()),
                        ('' if self.eq is None else ' eq: ' + self.eq.__repr__()),
                        ('' if self.hi is None else ' hi: ' + self.hi.__repr__()), ']'])

def insert(node, string):
    if len(string) == 0:
        return node

    head = string[0]
    tail = string[1:]
    if node is None:
        node = Node(head)

    if head < node.char:
        node.lo = insert(node.lo, string)
    elif head > node.char:
        node.hi = insert(node.hi, string)
    else:
        if len(tail) == 0:
            node.endpoint = True
        else:
            node.eq = insert(node.eq, tail)

    return node

def search(node, string):
    if node is None or len(string) == 0:
        return False

    head = string[0]
    tail = string[1:]
    if head < node.char:
        return search(node.lo, string)
    elif head > node.char:
        return search(node.hi, string)
    else:
        # use 'and' for matches on complete words only,
        # versus 'or' for matches on string prefixes
        if len(tail) == 0 or node.endpoint:
            return True
        return search(node.eq, tail)

def suffixes(node):
    if node is not None:
        if node.endpoint:
            yield node.char

        if node.lo:
            for s in suffixes(node.lo):
                yield s
        if node.hi:
            for s in suffixes(node.hi):
                yield s
        if node.eq:
            for s in suffixes(node.eq):
                yield node.char + s

def autocompletes(node, string):
    if node is None or len(string) == 0:
        return []

    head = string[0]
    tail = string[1:]
    if head < node.char:
        return autocompletes(node.lo, string)
    elif head > node.char:
        return autocompletes(node.hi, string)
    else:
        if len(tail) == 0:
            # found the node containing the prefix string,
            # so get all the possible suffixes underneath
            return suffixes(node.eq)
        return autocompletes(node.eq, string[1:])

class Trie:
    # a simple wrapper
    root = None
    def __init__(self, string):
        self.append(string)
    def append(self, string):
        self.root = insert(self.root, string)
    def __contains__(self, string):
        return search(self.root, string)
    def autocomplete(self, string):
        return map(lambda x: string + x, autocompletes(self.root, string))


trie = Trie('cute')
trie.append('cup')
trie.append('at')
trie.append('as')
trie.append('he')
trie.append('us')
trie.append('i')
print(trie.root)
print(list(trie.autocomplete('cu')))