首页 > 代码库 > 字典树用于单词联想

字典树用于单词联想

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

贴上代码

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
<style>p.p1 { margin: 0.0px 0.0px 0.0px 0.0px; font: 14.0px Courier } span.s1 { }</style>

>> 1

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

>> [‘abc‘, ‘alter‘]

>> abc

>> alter

>> hello

>> salt

 

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

字典树用于单词联想