Trie树实现词频统计与查找

#encoding:utf-8
from collections import defaultdict
import sys
reload(sys)
sys.setdefaultencoding('utf8')
class LBTrie:
"""
simple implemention of Trie in Python.
"""
def __init__(self):
self.trie = {}
self.size = 0

#添加单词
def add(self, word):
p = self.trie
dicnum = 0
word = word.strip()
for c in word:
if not c in p:
p[c] = {}
dicnum+=1
p = p[c]

if word != '':
#在单词末尾处添加键值''作为标记，即只要某个字符的字典中含有''键即为单词结尾
p[''] = ''
if dicnum == len(word):
return True
#查询单词
def search(self, word):
p = self.trie
word = word.lstrip()
for c in word:
if not c in p:
return False
p = p[c]
#判断单词结束标记''
if '' in p:
return True
return False

#打印Trie树的接口
def output(self):
#print '{'
self.__print_item(self.trie)
#print '}'
return  self.__print_item(self.trie)

#实现Trie树打印的私有递归函数，indent控制缩进
def __print_item(self, p, indent=0):
if p:
ind = '' + '\t' * indent
for key in p.keys():
label = "'%s' : " % key
print ind + label + '{'
self.__print_item(p[key], indent+1)

print ind + ' '*len(label) + '}'

def codeutil(strs):
return strs.decode('utf8','ignore').encode('GBK','ignore').decode('GBK','ignore')

if __name__ == '__main__':
trie_obj = LBTrie()
#添加单词
corpus = open('content.txt','r')
tree = open('tree.txt','w+')
countdic = defaultdict(int)
for record in corpus.readlines():
recordlist = record.split(' ')
for word in recordlist:
check = trie_obj.add(codeutil(word))
if check:
countdic[word] += 1
resortedcountdic = sorted(countdic.items(), key=lambda item: item[1], reverse=True)
for tup in resortedcountdic:
tree.write(''.join(codeutil(tup[0]))+'\t'+str(tup[1])+'\t')
#查找单词
if trie_obj.search(codeutil('氨基酸')):
print 'Yes'
else:
print 'No'