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

1. 字典树（Trie）

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

图-1 字典树

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

2. 字典树的代码实现

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

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

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

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

    def _add_child(self, char, value, overwrite=False):
        child = self._children.get(char)
        if child is None:
            child = Node(value)
            self._children[char] = child
        elif overwrite:
            child._value = value

        return child
1
2
3
4
5
6
7
8
9
10
11
12
13
14

上述代码中，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 None

    def __getitem__(self, key):
        state = self
        for char in key:
            state = state._children.get(char)
            if state is None:
                return None
        return state._value

    def __setitem__(self, key, value):
        state = self
        for 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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

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

3. 双数组字典树（DATrie）

用字典树这种数据结构查词时，最多比较$logN$次，比朴素方法效率高很多。不过，上述Trie类的代码中有一行语句效率不高，它就是state = state._children.get(char)。_children是一个字典，通过查找key来获取value。python中，key的查找是用散列法（hash），散列法本身效率很高，时间复杂度为$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 基值的选取

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

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

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

    def _add_child(self, char, value, overwrite=False):
        child = self._children.get(char)
        if child is None:
            child = Node(value)
            self._children[char] = child
        elif overwrite:
            child._value = value

        return child

    def get_children(self):
        return self._children

    def get_value_member(self):
        return self._value


class Trie(Node):
    def __init__(self):
        super().__init__(None)

    def __contains__(self, key):
        return self[key] is not None

    def __getitem__(self, key):
        state = self
        for char in key:
            state = state._children.get(char)
            if state is None:
                return None
        return state._value

    def __setitem__(self, key, value):
        state = self
        for 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] = val

        self.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) + 1
            self._check[p] = self._base[b]
            self._char[p] = ch

        for ch in keys:
            if ch == '\0':
                idx = children[ch].get_value_member()
                self._base[self._base[b]] = -idx - 1
                continue

            child_b = self._base[b] + ord(ch) + 1
            ch_children = children[ch].get_children()
            self.build(child_b, ch_children)

    def get_char_member(self):
        return self._char

    def get_base_member(self):
        return self._base

    def get_check_member(self):
        return self._check

    def get_base(self, keys):
        b = 0
        while True:
            conflict = False
            for ch in keys:
                p = self._base_pool[b] + ord(ch) + 1
                if self._check[p] != 0:
                    b += 1
                    conflict = True
                    break

            if not conflict:
                base = self._base_pool.pop(b)
                break

        return base

    def retrieve(self, word):
        b = 0
        for w in word:
            b = self.transition(b, w)
            if b is None:
                break

        val = -1
        if b is not None:
            p = self._base[b]
            if self._check[p] == self._base[b]:
                val = -self._base[p] - 1

        return val

    def transition(self, b, ch):
        p = self._base[b] + ord(ch) + 1
        if self._check[p] == self._base[b]:
            return p

        return None
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

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

5. 用双数组字典树分词

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

words = ['下雪', '下雪天', '适合', '适度', '火锅', '下雨', '比较']
dic = {}
for i, w in enumerate(words):
    dic[w] = i
    
dat = DAT(dic)
1
2
3
4
5
6

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

    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 - 1
                    break
                s = sent[j:k+1]
                if self.retrieve(s) > 0:
                    word_list.append(s)
                    k = j - 1
                    break

        word_list.reverse()
        return word_list
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

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

写主调用函数：

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

    dat = DAT(dic)
    words = dat.split("下雪天比较适合吃火锅")
    print(words)
1
2
3
4
5
6
7
8
9

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