# James Tauber
# http://jtauber.com/ 

class Trie:
    """
    A Trie is like a dictionary in that it maps keys to values. However,
    because of the way keys are stored, it allows look up based on the
    longest prefix that matches.
    """

    def __init__(self):
        self.root = [None, {}]


    def add(self, key, value):
        """
        Add the given value for the given key.
        """
        
        curr_node = self.root
        for ch in key:
            curr_node = curr_node[1].setdefault(ch, [None, {}])
        curr_node[0] = value


    def find(self, key):
        """
        Return the value for the given key or None if key not found.
        """
        
        curr_node = self.root
        for ch in key:
            try:
                curr_node = curr_node[1][ch]
            except KeyError:
                return None
        return curr_node[0]


    def find_prefix(self, key):
        """
        Find as much of the key as one can, by using the longest
        prefix that has a value. Return (value, remainder) where
        remainder is the rest of the given string.
        """
        
        curr_node = self.root
        remainder = key
        for ch in key:
            try:
                curr_node = curr_node[1][ch]
            except KeyError:
                return (curr_node[0], remainder)
            remainder = remainder[1:]
        return (curr_node[0], remainder)


    def convert(self, keystring):
        """
        convert the given string using successive prefix look-ups.
        """
        
        valuestring = ""
        key = keystring
        while key:
            value, key = self.find_prefix(key)
            if not value:
                return (valuestring, key)
            valuestring += value
        return (valuestring, key)
        
        
if __name__ == "__main__":    
    t = Trie()
    t.add("foo", "A")
    t.add("fo", "B")
    t.add("l", "C")

    assert t.find("fo") == "B"
    assert t.find("fool") == None

    assert t.find_prefix("fo") == ("B", "")
    assert t.find_prefix("fool") == ("A", "l")

    assert t.convert("fo") == ("B", "")
    assert t.convert("fool") == ("AC", "")