#!/usr/bin/env python
# encoding: utf-8
"""
@date:    20131001
@version: 0.2
@author:  wklken@yeah.net
@desc:    搜索下拉提示，基于后台提供数据，建立数据结构(前缀树)，用户输入query前缀时，可以提示对应query前缀补全@update:20131001 基本结构，新增，搜索等基本功能20131005 增加缓存功能，当缓存打开，用户搜索某个前缀超过一定次数时，进行缓存，减少搜索时间20140309 修改代码，降低内存占用@TODO:test case加入拼音的话，导致内存占用翻倍增长，要考虑下如何优化节点，共用内存"""
#这是实现cache的一种方式，也可以使用redis/memcached在外部做缓存
#https://github.com/wklken/suggestion/blob/master/easymap/suggest.py
#一旦打开，search时会对每个节点做cache，当增加删除节点时，其路径上的cache会被清除,搜索时间降低了一个数量级
#代价：内存消耗, 不需要时可以关闭,或者通过CACHED_THREHOLD调整缓存数量#开启
#CACHED = True
#关闭
CACHED = False#注意，CACHED_SIZE >= search中的limit，保证search从缓存能获取到足够多的结果
CACHED_SIZE = 10
#被搜索超过多少次后才加入缓存
CACHED_THREHOLD = 10############### start ######################class Node(dict):def __init__(self, key, is_leaf=False, weight=0, kwargs=None):"""@param key: 节点字符@param is_leaf: 是否叶子节点@param weight: 节点权重, 某个词最后一个字节点代表其权重，其余中间节点权重为0，无意义@param kwargs: 可传入其他任意参数，用于某些特殊用途"""self.key = keyself.is_leaf = is_leafself.weight = weight#缓存,存的是node指针self.cache = []#节点前缀搜索次数，可以用于搜索query数据分析self.search_count = 0#其他节点无关仅和内容相关的参数if kwargs:for key, value in kwargs.iteritems():setattr(self, key, value)def __str__(self):return '<Node key:%s is_leaf:%s weight:%s Subnodes: %s>' % (self.key, self.is_leaf, self.weight, self.items())def add_subnode(self, node):"""添加子节点:param node: 子节点对象"""self.update({node.key: node})def get_subnode(self, key):"""获取子节点:param key: 子节点key:return: Node对象"""return self.get(key)def has_subnode(self):"""判断是否存在子节点:return: bool"""return len(self) > 0def get_top_node(self, prefix):"""获取一个前缀的最后一个节点(补全所有后缀的顶部节点):param prefix: 字符转前缀:return: Node对象"""top = selffor k in prefix:top = top.get_subnode(k)if top is None:return Nonereturn topdef depth_walk(node):"""递归，深度优先遍历一个节点,返回每个节点所代表的key以及所有关键字节点(叶节点)@param node: Node对象"""result = []if node.is_leaf:#result.append(('', node))    if len(node) >0:#修改，避免该前缀刚好是关键字时搜索不到result.append((node.key[:-1], node))node.is_leaf=Falsedepth_walk(node)else:return [('', node)]if node.has_subnode():for k in node.iterkeys():s = depth_walk(node.get(k))#print k , s[0][0]result.extend([(k + subkey, snode) for subkey, snode in s])return result#else:#print node.key#return [('', node)]def search(node, prefix, limit=None, is_case_sensitive=False):"""搜索一个前缀下的所有单词列表 递归@param node: 根节点@param prefix: 前缀@param limit: 返回提示的数量@param is_case_sensitive: 是否大小写敏感@return: [(key, node)], 包含提示关键字和对应叶子节点的元组列表"""if not is_case_sensitive:prefix = prefix.lower()node = node.get_top_node(prefix)#print 'len(node):' ,len(node)#如果找不到前缀节点，代表匹配失败，返回空if node is None:return []#搜索次数递增node.search_count += 1if CACHED and node.cache:return node.cache[:limit] if limit is not None else node.cache#print depth_walk(node)result = [(prefix + subkey, pnode) for subkey, pnode in depth_walk(node)]result.sort(key=lambda x: x[1].weight, reverse=True)if CACHED and node.search_count >= CACHED_THREHOLD:node.cache = result[:CACHED_SIZE]#print len(result)return result[:limit] if limit is not None else result#TODO: 做成可以传递任意参数的，不需要每次都改    2013-10-13 done
def add(node, keyword, weight=0, **kwargs):"""加入一个单词到树@param node: 根节点@param keyword: 关键词，前缀@param weight: 权重@param kwargs: 其他任意存储属性"""one_node = nodeindex = 0last_index = len(keyword) - 1for c in keyword:if c not in one_node:if index != last_index:one_node.add_subnode(Node(c, weight=weight))else:one_node.add_subnode(Node(c, is_leaf=True, weight=weight, kwargs=kwargs))one_node = one_node.get_subnode(c) else:one_node = one_node.get_subnode(c)if CACHED:one_node.cache = []if index == last_index:one_node.is_leaf = Trueone_node.weight = weightfor key, value in kwargs:setattr(one_node, key, value)index += 1def delete(node, keyword, judge_leaf=False):"""从树中删除一个单词@param node: 根节点@param keyword: 关键词，前缀@param judge_leaf: 是否判定叶节点，递归用,外部调用使用默认值"""# 空关键词，传入参数有问题，或者递归调用到了根节点,直接返回if not keyword:returntop_node = node.get_top_node(keyword)if top_node is None:return#清理缓存if CACHED:top_node.cache = []#递归往上，遇到节点是某个关键词节点时，要退出if judge_leaf:if top_node.is_leaf:return#非递归，调用deleteelse:if not top_node.is_leaf:returnif top_node.has_subnode():#存在子节点，去除其标志 donetop_node.is_leaf = Falsereturnelse:#不存在子节点，逐层检查删除节点this_node = top_nodeprefix = keyword[:-1]top_node = node.get_top_node(prefix)del top_node[this_node.key]delete(node, prefix, judge_leaf=True)##############################
#  增补功能 读数据文件建立树 #
##############################def build(file_path, is_case_sensitive=False):"""从文件构建数据结构, 文件必须utf-8编码,可变更@param file_path: 数据文件路径，数据文件默认两列，格式“关键词\t权重"@param is_case_sensitive: 是否大小写敏感"""node = Node("")f = open(file_path)for line in f:line = line.strip()if not isinstance(line,unicode):line = line.decode('utf-8')parts = line.split('\t')name = parts[0]if not is_case_sensitive:name = name.lower()add(node, name, int(parts[1]))f.close()return nodeimport time
if __name__ == '__main__':#print '============ test1 ==============='#n = Node("")#default weight=0, 后面的参数可以任意加,搜索返回结果再从node中将放入对应的值取出,这里放入一个othervalue值#add(n, u'he', othervalue="v-he")#add(n, u'her', weight=0, othervalue="v-her")#add(n, u'hero', weight=10, othervalue="v-hero")#add(n, u'hera', weight=3, othervalue="v-hera")#delete(n, u'hero')#print "search h: "#for key, node in search(n, u'h'):#print key, node, node.othervalue, id(node)#print key, node.weight#print "serch her: "#for key, node in search(n, u'her'):#print key, node, node.othervalue, id(node)#print key, node.weightstart= time.clock()print '============ test2 ==============='tree = build("./shanxinpoi.txt", is_case_sensitive=False)print len(tree),'time:',time.clock()-startstartline=time.clock()print u'search 秦岭'for key, node in search(tree, u'秦岭', limit=10):print key, node.weightprint time.clock()-startline

# -*- coding:utf-8 -*-# base
# https://linux.thai.net/~thep/datrie/datrie.html
# http://jorbe.sinaapp.com/2014/05/11/datrie/
# http://www.hankcs.com/program/java/%E5%8F%8C%E6%95%B0%E7%BB%84trie%E6%A0%91doublearraytriejava%E5%AE%9E%E7%8E%B0.html
# (komiya-atsushi/darts-java | 先建立Trie树，再构造DAT，为siblings先找到合适的空间)
# https://blog.csdn.net/kissmile/article/details/47417277
# http://nark.cc/p/?p=1480
#https://github.com/midnight2104/midnight2104.github.io/blob/58b5664b3e16968dd24ac5b1b3f99dc21133b8c4/_posts/2018-8-8-%E5%8F%8C%E6%95%B0%E7%BB%84Trie%E6%A0%91(DoubleArrayTrie).md# 不需要构造真正的Trie树，直接用字符串，构造对应node，因为words是排过序的
# todo : error info
# todo : performance test
# todo : resize
# warning: code=0表示叶子节点可能会有隐患(正常词汇的情况下是ok的)
# 修正: 由于想要回溯字符串的效果，叶子节点和base不能重合(这样叶子节点可以继续记录其他值比如频率)，叶子节点code: 0->-1
# 但是如此的话，叶子节点可能会与正常节点冲突？ 找begin的使用应该是考虑到的？
#from __future__ import print_function
class DATrie(object):class Node(object):def __init__(self, code, depth, left, right):self.code = codeself.depth = depthself.left = leftself.right = rightdef __init__(self):self.MAX_SIZE = 2097152  # 65536 * 32self.base = [0] * self.MAX_SIZEself.check = [-1] * self.MAX_SIZE  # -1 表示空self.used = [False] * self.MAX_SIZEself.nextCheckPos = 0  # 详细 见后面->当数组某段使用率达到某个值时记录下可用点，以便下次不再使用self.size = 0  # 记录总共用到的空间# 需要改变size的时候调用，这里只能用于build之前。cuz没有打算复制数据.def resize(self, size):self.MAX_SIZE = sizeself.base = [0] * self.MAX_SIZEself.check = [-1] * self.MAX_SIZEself.used = [False] * self.MAX_SIZE# 先决条件是self.words ordered 且没有重复# siblings至少会有一个def fetch(self, parent):   ###获取parent的孩子，存放在siblings中，并记录下其左右截至depth = parent.depthsiblings = []  # size == parent.right-parent.lefti = parent.left while i < parent.right: #遍历所有子节点，right-left+1个单词s = self.words[i][depth:]  #词的后半部分if s == '':siblings.append(self.Node(code=-1, depth=depth+1, left=i, right=i+1)) # 叶子节点else:c = ord(s[0])  #字符串中每个汉字占用3个字符（code,实际也就当成符码），将每个字符转为数字 ，树实际是用这些数字构建的#print type(s[0]),cif siblings == [] or siblings[-1].code != c:siblings.append(self.Node(code=c, depth=depth+1, left=i, right=i+1)) # 新建节点else:  # siblings[-1].code == csiblings[-1].right += 1   #已经是排过序的可以直接计数+1i += 1# siblingsreturn siblings# 在insert之前，认为可以先排序词汇，对base的分配检查应该是有利的# 先构建树，再构建DAT，再销毁树def build(self, words):words = sorted(list(set(words)))  # 去重排序#for word in words:print word.decode('utf-8')self.words = words# todo: 销毁_root_root = self.Node(code=0, depth=0, left=0, right=len(self.words))  #增加第一个节点self.base[0] = 1siblings = self.fetch(_root)#for ii in  words: print ii.decode('utf-8')#print 'siblings len',len(siblings)#for i in siblings: print i.codeself.insert(siblings, 0)  #插入根节点的第一层孩子# while False:  # 利用队列来实现非递归构造# passdel self.wordsprint("DATrie builded.")def insert(self, siblings, parent_base_idx):""" parent_base_idx为父节点base index, siblings为其子节点们 """# 暂时按komiya-atsushi/darts-java的方案# 总的来讲是从0开始分配beigin]#self.used[parent_base_idx] = True
begin = 0pos = max(siblings[0].code + 1, self.nextCheckPos) - 1 #从第一个孩子的字符码位置开始找，因为排过序，前面的都已经使用nonzero_num = 0  # 非零统计first = 0  begin_ok_flag = False  # 找合适的beginwhile not begin_ok_flag:pos += 1if pos >= self.MAX_SIZE:raise Exception("no room, may be resize it.")if self.check[pos] != -1 or self.used[pos]:   # check——check数组，used——占用标记，表明pos位置已经占用nonzero_num += 1  # 已被使用continueelif first == 0:self.nextCheckPos = pos  # 第一个可以使用的位置，记录？仅执行一遍first = 1begin = pos - siblings[0].code  # 第一个孩子节点对应的beginif begin + siblings[-1].code >= self.MAX_SIZE:raise Exception("no room, may be resize it.")if self.used[begin]:    #该位置已经占用continueif len(siblings) == 1:  #只有一个节点begin_ok_flag = Truebreakfor sibling in siblings[1:]:if self.check[begin + sibling.code] == -1 and self.used[begin + sibling.code] is False: #对于sibling，begin位置可用begin_ok_flag = Trueelse:begin_ok_flag = False  #用一个不可用，则begin不可用break# 得到合适的begin# -- Simple heuristics --# if the percentage of non-empty contents in check between the# index 'next_check_pos' and 'check' is greater than some constant value# (e.g. 0.9), new 'next_check_pos' index is written by 'check'.#从位置 next_check_pos 开始到 pos 间，如果已占用的空间在95%以上，下次插入节点时，直接从 pos 位置处开始查找成功获得这一层节点的begin之后得到，影响下一次执行insert时的查找效率if (nonzero_num / (pos - self.nextCheckPos + 1)) >= 0.95:self.nextCheckPos = posself.used[begin] = True# base[begin] 记录 parent chr  -- 这样就可以从节点回溯得到字符串 # 想要可以回溯的话，就不能在字符串末尾节点记录值了，或者给叶子节点找个0以外的值？ 0->-1#self.base[begin] = parent_base_idx     #【*】#print 'begin:',begin,self.base[begin]if self.size < begin + siblings[-1].code + 1:self.size = begin + siblings[-1].code + 1for sibling in siblings: #更新所有子节点的check     base[s]+c=t & check[t]=sself.check[begin + sibling.code] = beginfor sibling in siblings:  # 由于是递归的情况，需要先处理完check# darts-java 还考虑到叶子节点有值的情况，暂时不考虑(需要记录的话，记录在叶子节点上)if sibling.code == -1:self.base[begin + sibling.code] = -1 * sibling.left - 1else:new_sibings = self.fetch(sibling)h = self.insert(new_sibings, begin + sibling.code) #插入孙子节点，begin + sibling.code为子节点的位置self.base[begin + sibling.code] = h #更新base所有子节点位置的转移基数为[其孩子最合适的begin]return begindef search(self, word):""" 查找单词是否存在 """p = 0  # rootif word == '':return Falsefor c in word:c = ord(c)next = abs(self.base[p]) + c# print(c, next, self.base[next], self.check[next])if next > self.MAX_SIZE:  # 一定不存在return False# print(self.base[self.base[p]])if self.check[next] != abs(self.base[p]):return Falsep = next# print('*'*10+'\n', 0, p, self.base[self.base[p]], self.check[self.base[p]])# 由于code=0,实际上是base[leaf_node->base+leaf_node.code]，这个负的值本身没什么用# 修正：left code = -1if self.base[self.base[p] - 1] < 0 and self.base[p] == self.check[self.base[p] - 1] :  #print preturn Trueelse:  # 不是词尾return Falsedef common_prefix_search(self, content):""" 公共前缀匹配 """# 用了 darts-java 写法，再仔细看一下result = []b = self.base[0]  # 从root开始p = 0n = 0tmp_str = ""for c in content:c = ord(c)p = bn = self.base[p - 1]      # for iden leafif b == self.check[p - 1] and n < 0:result.append(tmp_str)tmp_str += chr(c)#print(tmp_str )p = b + c   # cur nodeif b == self.check[p]:b = self.base[p]  # next baseelse:                 # no next nodereturn result# 判断最后一个nodep = bn = self.base[p - 1]if b == self.check[p - 1] and n < 0:result.append(tmp_str)return resultdef Find_Last_Base_index(self, word):b = self.base[0]  # 从root开始p = 0#n = 0#print len(word)tmp_str = ""for c in word:c = ord(c)p = bp = b + c   # cur node, p is new base position, b is the oldif b == self.check[p]:tmp_str += chr(c)b = self.base[p]  # next baseelse:                 # no next nodereturn -1#print '====', p, self.base[p], tmp_str.decode('utf-8')return pdef GetAllChildWord(self,index):result = []#result.append("")# print self.base[self.base[index]-1],'++++'if self.base[self.base[index]-1] <= 0 and self.base[index] == self.check[self.base[index] - 1]:result.append("")#return resultfor i in range(0,256):#print(chr(i))if self.check[self.base[index]+i]==self.base[index]:#print self.base[index],(chr(i)),ifor s in self.GetAllChildWord(self.base[index]+i):#print sresult.append( chr(i)+s)return resultdef FindAllWords(self, word):result = []last_index=self.Find_Last_Base_index(word)if last_index==-1:return resultfor end in self.GetAllChildWord(last_index):result.append(word+end)return resultdef get_string(self, chr_id):""" 从某个节点返回整个字符串, todo:改为私有 """if self.check[chr_id] == -1:raise Exception("不存在该字符。")child = chr_ids = []while 0 != child:base = self.check[child]print(base, child)label = chr(child - base)s.append(label)print(label)child = self.base[base]return "".join(s[::-1])def get_use_rate(self):""" 空间使用率 """return self.size / self.MAX_SIZEif __name__ == '__main__':words = ["一举","一举一动",'11',"一举成名","一举成名天下知","洛阳市西工区中州中路","人民东路2号","中州东","洛阳市","洛阳","洛神1","洛神赋","万科","万达3","万科翡翠","万达广场","洛川","洛川苹果","商洛","商洛市","商朝","商业","商业模","商业模式","万能","万能胶"]#for word in words:print [word]  #一个汉字的占用3个字符，words=[]for line in open('1000.txt').readlines():#    #print line.strip().decode('utf-8')
        words.append(line.strip())datrie = DATrie()datrie.build(words)#for line in open('1000.txt').readlines():#    print(datrie.search(line.strip()),end=' ')#print('-'*10)#print(datrie.search("景华路"))#print('-'*10)#print(datrie.search("景华路号"))# print('-'*10)#for item in datrie.common_prefix_search("商业模式"): print(item.decode('utf-8'))#for item in datrie.common_prefix_search("商业模式"):print item.decode('utf-8')# print(datrie.common_prefix_search("一举成名天下知"))#print(datrie.base[:1000])# print('-'*10)# print(datrie.get_string(21520))#index=datrie.Find_Last_Base_index("商业")#print(index),'-=-=-='#print datrie.search("商业"),datrie.search("商业"),datrie.search("商业模式")#print index, datrie.check[datrie.base[index]+230],datrie.base[index]for ii in  datrie.FindAllWords('中州中路'):print ii.decode('utf-8')#print(datrie.Find_Last_Base_index("一举")[2].decode('utf-8'))
#print()