#!/usr/bin/env python
class Node(object):
def __init__(self, value=None):
self.value = value
self.children = []
class Trie(object):
def __init__(self):
self.root = Node("Root")
def add(self, key):
self._add(self.root, key, 0)
def _add(self, node, key, index):
if len(key) == index: return
next_node = None
for n in node.children:
if n.value == key[index]:
next_node = n
break
if next_node is None:
next_node = Node(key[index])
node.children.append(next_node)
print "Adding", next_node.value, "into", node.value
self._add(next_node, key, index+1)
def find(self, key):
return self._find(self.root, key, 0)
def _find(self, node, key, index):
if len(key) == index: return node
found = False
for n in node.children:
if n.value == key[index]:
found = True
last_found_node = self._find(n, key, index+1)
if not found: return None
else: return last_found_node
def search(self, key):
result = []
last_node = self.find(key)
if last_node is None: return result
self._search(last_node, key, "", result)
for i in xrange(0, len(result)):
result[i] = key + result[i]
return result
def _search(self, node, key, char, result):
if len(node.children) == 0:
result.append(char)
return
for n in node.children:
self._search(n, key, char + n.value, result)
class WordCompletion(object):
def __init__(self, words):
self.trie = Trie()
for word in words:
print "===== Inserting", word, "====="
self.trie.add(word)
def search(self, prefix):
return self.trie.search(prefix)
if __name__ == "__main__":
words = ["ABCD", "ABCE", "AFG", "AFHIJ", "AFHIK", "XY"]
wc = WordCompletion(words)
assert ['ABCD', 'ABCE', 'AFG', 'AFHIJ', 'AFHIK'] == wc.search("A")
assert ['AFHIJ', "AFHIK"] == wc.search("AFHI")
assert ['AFHIJ', 'AFHIK'] == wc.search("AFH")
assert ['ABCD', 'ABCE'] == wc.search("ABC")
assert [] == wc.search("whatever")
Wednesday, July 11, 2012
Implementing Word Completion in Python
Subscribe to:
Post Comments (Atom)
No comments:
Post a Comment