字典树用于单词联想

最近要做一个单词联想的功能，经过调研选择使用字典树，节省空间，查找快。

贴上代码

class Trie(object):
    def __init__(self):
        self.path = {}
        self.value = None
        self.valid = False
        
    def __setitem__(self, key, val):
        head = key[0]
        if head in self.path:
            node = self.path[head]
        else:
            node = Trie()
            self.path[head] = node
            
        if len(key) > 1:
            remains = key[1:]
            node.__setitem__(remains, val)
        else:
            node.value = val
            node.valid = True
            
    def __getitem__(self, key):
        head = key[0]
        if head not in self.path:
            raise KeyError(key)
        else:
            node = self.path[head]
        
        if len(key) > 1:
            remains = key[1:]
            try:
                return node.__getitem__(remains)
            except KeyError:
                raise KeyError(key)
        elif node.valid:
            return node.value
        else:
            raise KeyError(key)
            
    def get(self, key, default=None):
        try:
            return self.__getitem__(key)
        except KeyError:
            return default
                    
    def __contains__(self, key):
        try:
            self.__getitem__(key)
            return True
        except KeyError:
            return False
        
    def __keys__(self, prefix=[], seen=[]):
        result = []
        if self.valid:
            is_str = True
            val = ''
            for k in seen:
                if type(k) != str or len(k) > 1:
                    is_str = False
                    break
                val += k
            if is_str:
                result.append(val)
            else:
                result.append(prefix)
                
        if len(prefix) > 0:
            head = prefix[0]
            prefix = prefix[1:]
            if head in self.path:
                nextpaths = [head]
            else:
                nextpaths = []            
        else:
            nextpaths = self.path.keys()
            
        for key in nextpaths:
            nextseen = []
            nextseen.extend(seen)
            nextseen.append(key)
            result += self.path[key].__keys__(prefix, nextseen)
            
        return result
                   
    def keys(self, prefix=[]):
        return self.__keys__(prefix)
     
    def __iter__(self):
        for item in self.keys():
            yield item
            
        raise StopIteration
        
        
trie = Trie()
trie['abc'] = 1
trie['salt'] = 1
trie['hello'] = 1
trie['alter'] = 1
print trie['abc']
print trie.keys()
print trie.keys('a') 
for key in trie:
    print key

>> 1

>> ['abc', 'alter', 'hello', 'salt']

>> ['abc', 'alter']

>> abc

>> alter

>> hello

>> salt

 

代码参考 https://github.com/bdimmick/python-trie/blob/master/trie.py