Wednesday, July 11, 2012

Implementing Word Completion in Python


#!/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")

No comments:

Post a Comment