Trie Python

count.py

#!/usr/bin/env python3
import sys

"""
Use a naive implementation of a Trie: http://en.wikipedia.org/wiki/Trie
"""
class Trie:

    def __init__(self):
        self.children = {} #letter<->sub-Trie mappings
        self.count = 0

    def add(self, word):
        if len(word) == 0:
            self.count += 1
        else:
            if word[0] not in self.children:
                self.children[word[0]] = Trie()
            self.children[word[0]].add(word[1:])

    def get_count(self, word):
        if len(word) == 0:
            return self.count
        if word[0] in self.children:
            return self.children[word[0]].get_count(word[1:])
        return 0

    def print(self, limit=5,lvl=0):
        space = lvl*2*' '
        print(space,'[',self.count,']')
        for key,trie in self.children.items():
            print(space, key)
            trie.print(limit-1,lvl+1)

    def freqs(self):
        if self.count > 0:
            yield '', self.count
        for key,trie in self.children.items():
            for s,count in trie.freqs():
                yield key+s, count

if __name__ == '__main__':
    c = 0
    t = Trie()
    for word in open("passwords.txt"):
        t.add(word.strip())

    for s,c in t.freqs():
        print(s,c)

test.py

import unittest

from count import Trie 

class TestTrie(unittest.TestCase):

    def test_freqs(self):
        words = "ab,a,abc,ba,cda,aa,aa,aa".split(',')
        t = Trie()
        for word in words:
            t.add(word)
        freqs_list = list(t.freqs())
        freqs = {k:v for k,v in t.freqs()}
        self.assertEqual(freqs['aa'],3)
        self.assertEqual(freqs['ab'],1)
        self.assertEqual(len(freqs),len(set(words)))

if __name__ == '__main__':
    unittest.main()