使用python代码实现三叉搜索树高效率”自动输入提示”功能

使用python代码实现三叉搜索树高效率”自动输入提示”功能

 

原始出处： 韩少杰| ThinkDevil

最近项目中需要通过全拼音和简写拼音实现输入自动提示结果功能，查了一些资料发现三叉搜索树无论是在时间还是空间上都比较优秀。

三叉搜索树是trie树的演化版，除去了指针，这样在空间上节省不少，每个节点基本分为三个方向：左、中、右，当字符小于当前节点则存放左边，大于则存放右边，等于则存放中间。

具体实现原理可参考：http://igoro.com/archive/efficient-auto-complete-with-a-ternary-search-tree/

 

假如我们存入”AB”,”ABBA”.”ABCD”.”BCD”,这几个单词，那么三叉树就会出现如下的储存方式：

 

虚线表示匹配的箭头，黄色的表示单词结尾

 

下面是python实现代码:

 

# -*- coding: utf-8 -*-

#Ternary Search Tree

 

# 小于的left, 大于的right, 等于的mid

 

#定义状态属性

_SENTINEL = ()

 

class TST(object):

    #定义三叉树位置结构

    __slots__ = ('splitchar','l','m','r','v')

 

    #初始化结构

    def __init__(self, ch=None):

        self.splitchar = ch

        self.l = self.m = self.r = None

 

    #返回状态

    def __getstate__(self):

        l = [self.splitchar,self.l,self.m,self.r]

        if hasattr(self,'v'):

            l.append(self.v)

        return tuple(l)

 

    #设置状态,目的支持pickle的持久化对象

    def __setstate__(self,l):

        self.splitchar = l[0]

        self.l = l[1]

        self.m = l[2]

        self.r = l[3]

        if len(l) > 4 :

            self.v = l[4]

 

    #定义类方法，递归方式插入字母,这样不用实例

    @classmethod

    def insert(klass,p,k,v):

        #获取字母

        ch = k[0]

 

        #若三叉树结构为空，则初始化

        if p is None:

            p = TST(ch)

        elif p.splitchar is None:

            p.splitchar = ch

 

        #若当前字符小于节点字符，则做插入

        if ch < p.splitchar:

            p.l = klass.insert(p.l,k,v)

        #若当前字符等于节点字符，则

        elif ch == p.splitchar:

            #获取剩余字符

            k = k[1:]

            if k:

                p.m = klass.insert(p.m,k,v)

            else:

                #标记字母位置

                p.v = v

        #否则右插入

        else:

            p.r = klass.insert(p.r,k,v)

 

        return p

 

    #添加数据

    def add(self,k,v):

        return self.insert(self,k,v)

 

    #搜索字符串

    def search(self,s,fallback=None):

        p = self

        while p:

            ch = s[0]

            if ch < p.splitchar:

                p = p.l

            elif ch == p.splitchar:

                s = s[1:]

                if not s:

                    if hasattr(p,'v') :

                        return p.v

                    break

                p = p.m

            else:

                p = p.r

        return fallback

 

    #搜索前缀的字符

    def prefix_search(self,s):

        p = self

        while p:

            ch = s[0]

            if ch < p.splitchar:

                p = p.l

            elif ch == p.splitchar:

                s = s[1:]

                if not s:

                    return list(p)

                p = p.m

            else:

                p = p.r

        return []

 

    #批量增加数据

    def bulk_add(self,l,start=0,stop=None,sorted=False):

        '''

        为了取得最佳性能，字符串应该以随机或者自平衡顺序插入到三叉树中

        尤其不能按照字母顺序，这样就和链表没啥区别了

        '''

        #若是没有排序的数据则进行排序

        if not sorted:

            l.sort()

 

        #若没有结束位置，则以全部长度作为结束

        if stop is None:

            stop = len(l)

 

        #比较开始到结束距离

        diff = stop - start

 

        #若为一个则直接添加

        if diff == 1 :

            self.add(l[start][0],l[start][1])

        #若为两个同样直接添加

        elif diff == 2 :

            self.add(l[start][0],l[start][1])

            self.add(l[start+1][0],l[start+1][1])

            return

        #两个以上则开始计算中间值

        else:

            mid_p = start + (diff / 2)

            #增加中间值

            self.add(l[mid_p][0],l[mid_p][1])

            #采用分治法递归增加，让我回忆起快速排序

            self.bulk_add(l,mid_p+1,stop,True)

            self.bulk_add(l,start,mid_p,True)

 

    def __contains__(self,k):

        if self.search(k,_SENTINEL) is _SENTINEL:

            return False

        return True

 

    def __iter__(self):

        stack = []

        p = self

        if not p:

            return

        while True:

            if p.r:

                stack.append(p.r)

            if p.m:

                stack.append(p.m)

            if p.l:

                stack.append(p.l)

            if hasattr(p,'v') :

                yield p.v

            if not stack:

                break

            p = stack.pop()