Skip to content

Instantly share code, notes, and snippets.

@dpapathanasiou dpapathanasiou/trie.py
Last active Feb 23, 2019

Embed
What would you like to do?
Ternary Search Tree in python
#!/usr/bin/env python
"""
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))
@dpapathanasiou

This comment has been minimized.

Copy link
Owner Author

commented Nov 26, 2017

Usage, using the example on wikipedia:

>>> trie = Trie('cute')
>>> trie.append('cup')
>>> trie.append('at')
>>> trie.append('as')
>>> trie.append('he')
>>> trie.append('us')
>>> trie.append('i')
>>> trie.root
[c lo: [a eq: [t <end> lo: [s <end>]]] eq: [u eq: [t lo: [p <end>] eq: [e <end>]]] hi: [h eq: [e <end>] hi: [u lo: [i <end>] eq: [s <end>]]]]
>>> 'cup' in trie
True
>>> 'cups' in trie
False
>>> 'cute' in trie
True
@dpapathanasiou

This comment has been minimized.

Copy link
Owner Author

commented Dec 10, 2017

Now with an autocomplete function:

>>> trie.autocomplete("cu")
['cup', 'cute']
>>> trie.autocomplete("a")
['at', 'as']
>>> trie.autocomplete("mo")
[]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.