trie.py

class Trie(object):
    class Node(object):
        def __init__(self):
            self.val = 0
            self.children = {}

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.head = Trie.Node()


    def __getitem__(self, key):
        return self._get(key)

    def __contains__(self, key):
        try:
            self._get(key)
            return True
        except KeyError:
            return False

    def _get(self, key):
        curr = self.head
        for c in key:
            if c not in curr.children:
                raise KeyError(key)

            curr = curr.children[c]

        return curr.val

    def insert(self, key, val):
        """
        :type key: str
        :type val: int
        :rtype: void
        """
        curr = self.head

        for c in key:
            if c not in curr.children:
                curr.children[c] = Trie.Node()

            curr = curr.children[c]

        curr.val = val

    def sum(self, prefix):
        """
        :type prefix: str
        :rtype: int
        """
        total = 0
        curr = self.head

        for c in prefix:
            if c in curr.children:
                curr = curr.children[c]
            else:
               return 0

        return self._recur_sum(curr)

    def descendents(self, prefix):
        curr = self.head
        for c in prefix:
            if c in curr.children:
                curr = curr.children[c]
            else:
                return
                yield

        yield from self._descendents(curr)

    def _descendents(self, node):
        if not node.children:
            yield node.val
        else:
            for n in node.children.values():
                yield from self._descendents(n)

            yield node.val

    def _recur_sum(self, node):
        if not node.children:
            return node.val

        return node.val + sum([self._recur_sum(n) for n in node.children.values()])


if __name__ == '__main__':
    t = Trie()
    t.insert("hello", 100)
    t.insert("help", 1)
    t.insert("hadoop", 1000000)
    t.insert("show", 99)

    print(list(t.descendents("he")))