68 lines
1.7 KiB
Python
68 lines
1.7 KiB
Python
#!/usr/bin/env python
|
|
|
|
"""
|
|
Quick & dirty prefix tree (aka trie).
|
|
"""
|
|
|
|
class Trie(object):
|
|
def __init__(self):
|
|
self.complete = False # Does this node complete a valid word?
|
|
self.children = {}
|
|
|
|
def add(self, value):
|
|
if not value:
|
|
self.complete = True
|
|
return
|
|
prefix = value[0]
|
|
remainder = value[1:]
|
|
node = self.children.get(prefix, None)
|
|
if not node:
|
|
node = Trie()
|
|
self.children[prefix] = node
|
|
node.add(remainder)
|
|
|
|
def find(self, value):
|
|
"Return the node associated with a value (None if not found)."
|
|
if not value:
|
|
return self
|
|
node = self.children.get(value[0], None)
|
|
if not node:
|
|
return None
|
|
return node.find(value[1:])
|
|
|
|
def dump(self, indent=0):
|
|
for key in sorted(self.children.keys()):
|
|
text = indent * ' '
|
|
text += key
|
|
node = self.children[key]
|
|
if node.complete:
|
|
text += '*'
|
|
print(text)
|
|
node.dump(indent+2)
|
|
|
|
if __name__ == '__main__':
|
|
t = Trie()
|
|
t.add('hell')
|
|
t.add('hello')
|
|
t.add('he')
|
|
t.add('world')
|
|
t.add('alphabet')
|
|
t.add('foo')
|
|
t.add('food')
|
|
t.add('foodie')
|
|
t.add('bar')
|
|
t.add('alphanumeric')
|
|
t.dump()
|
|
|
|
assert t.find('r') is None
|
|
assert t.find('bars') is None
|
|
assert not t.find('hel').complete
|
|
assert t.find('hell').complete
|
|
assert t.find('hello').complete
|
|
assert not t.find('f').complete
|
|
assert not t.find('fo').complete
|
|
assert t.find('foo').complete
|
|
assert t.find('food').complete
|
|
assert not t.find('foodi').complete
|
|
assert t.find('foodie').complete
|