用Python实现字典树(Trie)与双数组字典树(DATrie)

1. 字典树（Trie）

假如我们把字典中的词以记录的形式（无序）存入数据库中。现给定一串字符，要查找该字符串是否为字典中的词。因为数据库中的记录是无序的，所以，最朴素的方法就逐记录匹配。此方法简单，但是效率不高。因为每次匹配都必须从首字符开始。当然，可以将数据库中的记录按字符编码进行排序。这样，首字相同的词的记录都聚集在某个区间，匹配首字时直接跳至首字所处的区间开头即可，其它字符的匹配以此类推。其实这种方法的思想就是前缀匹配的思想。但是用数据库实现比较麻烦，可以用字典树这种数据结构来实现。
先来一张示意图，直观感受一下字典树的结构：

图-1 字典树

例如，要查找“下雪天”是否在字典中，先检查根结点是否存在孩子结点：“下”。如果存在，接下来只需在以”下“为根结点的子树中搜索即可，极大地缩小了检索范围。在自上而下匹配过程中，如果某个字匹配不到结点，说明该词不在字典中，立即结束检索过程。

2. 字典树的代码实现

虽然Python没有指针概念，不过，我们可以用字典来表示结点所拥有的孩子结点。所以，结点类应该有一个成员变量children，它的数据类型是字典。我们再设置一个成员变量用来存储结点的值，代码如下所示：

class Node(object):def __init__(self, value):self._children = {}self._value = value

可以用字符作为字典元素的key，如图-1中圆圈里边的字。用孩子结点的地址作为字典元素的value。在Node类中，我们要实现一个添加孩子结点的方法，其主要功能就是往children字典中添加元素，代码如下：

class Node(object):def __init__(self, value):self._children = {}self._value = valuedef _add_child(self, char, value, overwrite=False):child = self._children.get(char)if child is None:child = Node(value)self._children[char] = childelif overwrite:child._value = valuereturn child

上述代码中，child=Node(value) 中的child的值就是结点对象的地址。根结点的children变量的值类似长这样：{‘火’: node_1, ‘下’: node_2, ‘适’: node_3}
结点类构建好了，接下来构建字典树类。字典树是由结点构成的，所以字典树类继承结点类。字典树类要有一个方法能够把词语添加到树中。设字典树对象名为trie，那么，也就是要实现这个操作：trie[‘下雪’] = 1，trie[‘下雪天’]=2。这种赋值操作可以用Python的魔法函数__setitem__( )实现。
字典树类还要有查寻功能，即查寻某词是否在字典中。这里我们用这种形式查寻：isExist = trie[‘下午’]，根据isExist的值来判断是否查找成功。字典树类的代码如下所示：

class Trie(Node):def __init__(self):super().__init__(None)def __contains__(self, key):return self[key] is not Nonedef __getitem__(self, key):state = selffor char in key:state = state._children.get(char)if state is None:return Nonereturn state._valuedef __setitem__(self, key, value):state = selffor i, char in enumerate(key):if i < len(key) - 1:state = state._add_child(char, None, False)else:state = state._add_child(char, value, True)

为了简化操作，上述代码中将非词尾结点的value设为None。这样就可以根据__getitem__( )的返回值来判断查寻的某个字符串是否在字典中。如果返回值是None，则该词不在字典中。反之，则该词是字典中的词。

3. 双数组字典树（DATrie）

用字典树这种数据结构查词时，最多比较 l o g N log N logN次，比朴素方法效率高很多。不过，上述Trie类的代码中有一行语句效率不高，它就是state = state._children.get(char)。_children是一个字典，通过查找key来获取value。python中，key的查找是用散列法（hash），散列法本身效率很高，时间复杂度为 O ( 1 ) O(1) O(1)。问题出在python会把字符的unicode码转为64位编码（假设是64位OS）进行散列。这不仅耗内存，而且跨度大的散列效率也不高。

算法工程师的探索永无止境。既然python字典的做法不能满足我们的需求，有没有办法绕过字典查寻操作，也能获取到孩子结点。办法还是有的，那就是用两个数组来映射字典树Trie。这样在查找时，只需要对数组元素进行比较，而数组元素的取值范围可以由我们自己来定。

3.1 双数组字典树的基本思想

用通俗的话说，双数组字典树就是用两个数组来表示字典树结点间的父子关系。这两个数组分别是基数组和验证数组。假设用base表示基数组。用check表示验证数组。
如果为字典树中的每个结点都赋一个基值，将这个基值存储在下标为b的base数组中，那么，当下面式(2)成立时，表明两个结点存在父子关系。

p = base[b] + code (1)
check[p] == base[b] (2)

其中，code表示字符C的编码，b为数组下标，它表示字符C的父结点（暂且假定C有父结点）的基值所在的数组元素下标（base数组）。
用一张示意图来感受一下式(1)、式(2)表达的意思。

图-2

上图中，b所指位置的元素是结点’下’的父结点的基值，'下’结点的基值应该储存在base数组的p位置。

3.2 对原始双数组的两个改进

改进(一)：
为了能在双数组中判断某字是否是词尾，我们需要对图-1的字典树进行改进，也就是要在词尾结点后增加一个特殊的孩子结点，该结点字符可以用’\0’表示，如图-3所示。

图-3

改进(二)：

在根据图-3构建双数组时，如果当前结点是非’\0’结点，则 p = base[b] + code+1。如果当前结点是’\0’结点，则 p = base[b] + code，因为’\0’的unicode=0，所以p = base[b] 。这样就可以用双数组来判断某字符是否为词尾了。

3.3 基值的选取

基值是用来判断结点是否存在父子关系的，所以必须让每个结点的基值都是唯一的，以免其它结点认错父亲。
为了避免在写check[p]时，因check[p]不空闲，而重新调整基值，我们采用广度优先来遍历字典树。对check数组，采用一次性插入一群子结点的策略（属于当前根结点的所有子结点）。也就是先尝试着选取一个基值base_value，判断用这个基值计算出来的所有的p=base_val + code，是否都满足check[p]是空闲的。如果满足，base_val就确定了。如果不满足，就选取大一点的base_value，再试。
为了有效利用内存空间，基值按从小到大的顺序选取。
为了避免基值重复，这里有个小技巧。我们可以按升序生成一个基值序列，然后按从小到大的顺序从序列池中取基值。如果选出的基值满足要求，就把该基值从序列池中移除。如果不满足要求，就往后取一个大的，依此类推。
对于’\0’结点的基值，可以用负数。

4. 双数组字典树的代码实现

class Node(object):def __init__(self, value):self._children = {}self._value = valuedef _add_child(self, char, value, overwrite=False):child = self._children.get(char)if child is None:child = Node(value)self._children[char] = childelif overwrite:child._value = valuereturn childdef get_children(self):return self._childrendef get_value_member(self):return self._valueclass Trie(Node):def __init__(self):super().__init__(None)def __contains__(self, key):return self[key] is not Nonedef __getitem__(self, key):state = selffor char in key:state = state._children.get(char)if state is None:return Nonereturn state._valuedef __setitem__(self, key, value):state = selffor char in key:state = state._add_child(char, None, False)state._add_child('\0', value, True)class DAT(object):def __init__(self, dic):self._base = [0 for i in range(100000)]self._check = [0 for i in range(100000)]self._char = ['-' for i in range(100000)]self._base_pool = [i for i in range(1, 10000)]self._trie = Trie()for key, val in dic.items():self._trie[key] = valself.build(0, self._trie.get_children())def build(self, b, children):keys = children.keys()keys = sorted(keys, key=lambda x: ord(x))self._base[b] = self.get_base(keys)for ch in keys:if ch == '\0':p = self._base[b]else:p = self._base[b] + ord(ch) + 1self._check[p] = self._base[b]self._char[p] = chfor ch in keys:if ch == '\0':idx = children[ch].get_value_member()self._base[self._base[b]] = -idx - 1continuechild_b = self._base[b] + ord(ch) + 1ch_children = children[ch].get_children()self.build(child_b, ch_children)def get_char_member(self):return self._chardef get_base_member(self):return self._basedef get_check_member(self):return self._checkdef get_base(self, keys):b = 0while True:conflict = Falsefor ch in keys:p = self._base_pool[b] + ord(ch) + 1if self._check[p] != 0:b += 1conflict = Truebreakif not conflict:base = self._base_pool.pop(b)breakreturn basedef retrieve(self, word):b = 0for w in word:b = self.transition(b, w)if b is None:breakval = -1if b is not None:p = self._base[b]if self._check[p] == self._base[b]:val = -self._base[p] - 1return valdef transition(self, b, ch):p = self._base[b] + ord(ch) + 1if self._check[p] == self._base[b]:return preturn None

上述代码中，retrieve( )函数的功能是判断字符串是否在字典中，也就是判断字符串是否是字典中的词。如果是词，返回其在字典中的索引。如果不是词，则返回-1。

5. 用双数组字典树分词

首先，用上面的代码，构建双数组字典树。

words = ['下雪', '下雪天', '适合', '适度', '火锅', '下雨', '比较']
dic = {}
for i, w in enumerate(words):dic[w] = idat = DAT(dic)

写一个分词函数，用逆向最长匹配对句子进行分词，代码如下：

    def split(self, sent):"""逆向最长匹配:param sent::return:"""word_list = []k = len(sent)while k > 0:for j in range(0, k+2):if j == k+1:# 当前字符不在字典中word_list.append(s)k = k - 1breaks = sent[j:k+1]if self.retrieve(s) > 0:word_list.append(s)k = j - 1breakword_list.reverse()return word_list

split( )是DAT类的成员方法。

写主调用函数：

if __name__ == '__main__':words = ['下雪', '下雪天', '适合', '适度', '火锅', '下雨', '比较']dic = {}for i, w in enumerate(words):dic[w] = idat = DAT(dic)words = dat.split("下雪天比较适合吃火锅")print(words)

“下雪天比较适合吃火锅”，分词结果如下：