Trie树/字典树的原理及实现[C/C++]

前言

我们学习过很多字符串查询的算法，暴搜，KMP、BM、RK等字符串匹配算法，这些都是在文本中去查找我们的模式串。我们在搜索引擎的输入栏中输入时，常常仅输入了前缀，下面就给我们列出了几个含有此前缀的搜索项，这难道也是我们的字符串匹配算法吗？通过对字典树Trie的学习，相信会对此有所理解。

引例：Google经典面试题

经典面试题:搜索提示(自动补全)

如何根据用户输入的关键字，从我们的语料库中提取出一些用户可能想要的关键字呈现给用户呢？

用一棵字典树去存储一些语料库，根据用户输入的前缀去给出我们相应的字符串。

字典树的原理与实现

定义

字典树，又称trie树，是一种树形的数据结构。可以用作词频统计，利用字符串的公共前缀来减少查询时间。一般用来查找某个字符串S是否在一个字符串集合中。

字典树的结构

假如我们用字典树存储了如下单词集合：to，a，tea，ted，test，ice

那么对应的字典树如下：

对于上图：

• 每个节点都有对应下标，便于访问
• 下标从0开始，root用0标识
• 节点不存储字符，边存储字符
• 黄色节点代表有字符串该节点处结束

字典树的操作

字典树一般提供两种操作，插入和删除，都只需要遍历一遍字符串即可。

字符串插入

字符串的插入就是对字符串的遍历。

插入算法实现如下:

• 定义当前结点为字典树根结点root, 遍历给定字符串s;
• 对于字符串第i个字符s[i]，查询当前结点是否有s[i] 这个子结点;
  • 如果不存在，则创建一个新结点;
  • 如果存在，不作处理;
• 更新当前结点为原当前结点的i号子结点;
• 遍历完毕字符串s后在当前结点打上一个标记，标识结尾结点;

以前面的例子为例，如下动画展示了字符串插入过程：

初始只有一个root节点，第一个插入字符串是“to”，root就扩展成了如下单链结构

插入“a”，root由于没有’a’子节点的存在，所以又创建了’a’

插入“tea”，由于有“t“这一公共前缀的存在，所以我们只需要创建额外的’e’,‘a’

对于“ted”和“test”则类似

插入“ice”

如此，就完成了我们Trie的构造，假设字符串长度为m，一共插入n次，每次插入时间复杂度为O(m).总的时间复杂度就是O(nm)

字符串查询

字符串的查询也是对字符串的遍历。

查询算法实现如下:

• 定义当前结点cur为字典树根结点root, 遍历给定字符串s;
• 对于字符串第i个字符s[i]，查询当前结点是否有s[i]这个子结点;
  • 如果不存在，则返回false;
  • 如果存在，不作处理;
• 更新当前结点为原当前结点的i号子结点;
• 遍历完毕字符串s后，对当前结点cur判断是否存在结尾标记，存在则返回true,否则返回
  false;

我们仍以前面的字典树为例，进行“test”的查询，先从root寻找‘t’，到达节点1，再在节点1寻找‘e’，到达节点4，再在节点4寻找’s‘，到达节点7，再在节点7寻找’t‘，到达节点8，至此字符串遍历结束，而8号节点有结束标记，所以”test“在字典树中

但是如果我们查询”tes“，我们发现到7号节点结束，但是7号节点没有结束标志，所以”tes“不在字典树中。

字典树的实现

对于我们的字符有时会是大小写字母，有时会是数字，有时二者兼用，但对于不同的情况我们可以选择更为适合的Trie实现方式，这里我们介绍两种：字符集数组法和字符集映射法。

字符集数组法

每个结点保存一个长度固定为字符集大小(例如26)的数组，以字符为下标，保存指向的结点下标（不存指针是因为64位平台下指针为8字节，而int只有4字节），空间复杂度为O(结点数*字符集大小)，查询的时间复杂度为O(单词长度)
适用于较小字符集，或者单词短、分布稠密的字典

对于字典树，我们要设计两个类：节点类和字典树类。

节点类结构设计

对于节点类，我们要存储三个信息：结尾标记、词频统计、子节点数组。

这里对于子节点数组特殊说明，子节点数组可以直接存储子节点的地址，这是较为简单且直观的实现方法

但是这里我把所有创建的节点都放在了字典树类里面的一个数组里，这样我们可以用在节点中存储子节点下标的方式来完成对子节点的间接索引，因为64位平台上int为4字节，而指针为8字节，这样节省空间。

代码定义如下：

template //字符下标映射，后面再讲
struct TrieNode
{
    TrieNode();//构造函数
    
    bool isexist(int idx);//子节点查询
    bool isword();//当前节点是否是单词结尾
    
    void addnode(int idx, int node_idx);//增加子节点，子节点下标为node_idx
    int getNode(int idx);//获取子节点的下标

    void setword();//设置为单词结尾
    void addcnt();//增加词频
    
    bool _isword;//单词结尾标记
    int _cnt;//词频
    vector _nodes;//子节点数组
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

1. _isword即单词结尾标记，bool类型变量，完成一个字符串的插入后，在对应的末尾节点标记为true。
2. 记从根节点到当前节点的字符串为s，则_cnt就是字典树中以s为前缀的字符串数目
3. _nodes即一段连续数组，存储了每个子节点的下标

节点的接口

字符映射

字符都有对应的ASCII值，显然不能直接用来作为下标，通常都要进行转换，我们选择把对应的映射方式封装成仿函数，然后把仿函数的选择作为我们的模板参数，这样就实现了泛化。

小写字母下标映射

typedef struct HashLower // 小写字母映射
{
    int operator()(char ch)
    {
        return ch - 'a';
    }
    static const int _capacity;
} Lower;
1
2
3
4
5
6
7
8

大写字母下标映射

typedef struct HashUpper // 大写字母映射
{
    int operator()(char ch)
    {
        return ch - 'A';
    }
    static const int _capacity;
} Upper;
1
2
3
4
5
6
7
8

大小写混合字母下标映射

typedef struct Hash_U_L // 大小写字母混合映射
{
    int operator()(char ch)
    {
        if (ch >= 'a' && ch <= 'z')
            return ch - 'a';
        return ch - 'A' + 26;
    }
    static const int _capacity;
} U_L;
1
2
3
4
5
6
7
8
9
10

数字字符下标映射

typedef struct HashDigit // 数字字符映射
{
    int operator()(char ch)
    {
        return ch ^ 48;
    }
    static const int _capacity;
} Digit;
1
2
3
4
5
6
7
8

映射容量的设置
我们发现几个仿函数内都存放了_capacity这样一个静态整型常量，这是出于大小写字母各有26个，而大小写混合字母一共有52个，数字字符只有10个，所以不同的字符选择对应了不同的容量，而在不同的仿函数内通过设置静态变量存储容量可以很好的解决此问题，类方法的可维护性也很强。

const int HashLower::_capacity = 26;
const int HashUpper::_capacity = 26;
const int Hash_U_L::_capacity = 52;
const int HashDigit::_capacity = 10;
1
2
3
4

节点类的代码实现

template 
struct TrieNode
{
    TrieNode() : _isword(false), _cnt(0), _nodes(HashFunc::_capacity, -1)
    {
    }

    ~TrieNode()
    {
    }
    bool isexist(int idx)
    {
        return _nodes[idx] != -1;
    }
    bool isword()
    {
        return _isword;
    }
    void addnode(int idx, int node_idx)
    {
        _nodes[idx] = node_idx;
    }
    int getNode(int idx)
    {
        return _nodes[idx];
    }
    void setword()
    {
        _isword = true;
    }
    void addcnt()
    {
        _cnt++;
    }
    bool _isword;
    int _cnt;
    vector _nodes;
};
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

字典树类结构设计

template 
class Trie
{
private:
    HashFunc hash_id;//采用的字符映射
    typedef TrieNode Node;

public:
    Trie() : _root(1, Node())//初始化根节点
    {
    }

    void insert(const string &str);//插入字符串
    bool search(const string &str);//字符串查询
    
    Node *root();//对根节点封装
    Node *to_Node(int node_idx);//获取下标为node_idx的节点
    Node *genNode(int curidx);//产生新结点，返回curidx下标的节点地址

private:
    vector _root;
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

字典树接口实现

插入

前面已经进行过插入算法描述，直接看代码。

遍历字符串，子节点不存在对应字符就创建节点，当前节点移动到子节点，增加词频

这里解释下为什么cur在每次创建节点后都要重新定位，因为vector是一个动态增容的容器，底层用的是T*去存储，每次增容都会析构原有元素，再重新创建，导致我们的cur可能失效，所以要重定位。（其实可以直接用Node*代替vector，注意增容就行了，这里我强迫症，喜欢用vector，其实不太好）

  void insert(const string &str)
    {
        int idx, curidx = 0;
        Node *cur = root();
        for (auto ch : str)
        {
            idx = hash_id(ch);
            if (!cur->isexist(idx))
            {
                cur = genNode(curidx);
                cur->addnode(idx, _root.size() - 1);
            }
            curidx = cur->getNode(idx);
            cur = to_Node(curidx);
            cur->addcnt();
        }
        cur->setword();
    }
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

查询

遍历字符串，判断存不存在即可

    bool search(const string &str)
    {
        int idx;
        Node *cur = root();
        for (auto ch : str)
        {
            idx = hash_id(ch);
            if (!cur->isexist(idx))
            {
                return false;
            }
            cur = to_Node(cur->getNode(idx));
        }
        return cur->isword();
    }
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

其他简单接口直接给出

    Node *root()
    {
        return &_root[0];
    }
    Node *to_Node(int node_idx)
    {
        return &_root[node_idx];
    }
    Node *genNode(int curidx)
    {
        _root.emplace_back(Node());
        return &_root[curidx];
    }
1
2
3
4
5
6
7
8
9
10
11
12
13

字符集映射法（适用性广）

前面的字符集数组法，虽然也进行了泛化处理，但是适用性还是有些局限，把每个结点上的字符集数组改为一个映射(unordered_map),空间复杂度为O(文本字符总数)，查询的时间复杂度为O(单词长度)，但常数稍大一些，适用性更广

节点类结构设计

struct TrieNode
{
    TrieNode();//构造函数
    
    bool isexist(int idx);//子节点查询
    bool isword();//当前节点是否是单词结尾
    
    void addnode(int idx, int node_idx);//增加子节点，子节点下标为node_idx
    int getNode(int idx);//获取子节点的下标

    void setword();//设置为单词结尾
    void addcnt();//增加词频
    
    bool _isword;//单词结尾标记
    int _cnt;//词频
    unordered_map _nodes;//子节点映射
};
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

_nodes从数变为了映射，这样就不需要HashFunc来对字符下标进行映射了，也不需要单独处理容量

节点类的代码实现

由于代码基本和前面如出一辙所以直接给出代码

    struct TrieNode
    {
        TrieNode() : _isword(false), _cnt(0)
        {
        }

        ~TrieNode()
        {
        }
        bool isexist(char ch)
        {
            return _nodes.count(ch);
        }
        bool isword()
        {
            return _isword;
        }
        void addnode(char ch, int node_idx)
        {
            _nodes[ch] = node_idx;
        }
        int getNode(char ch)
        {
            return _nodes[ch];
        }
        void setword()
        {
            _isword = true;
        }
        void addcnt()
        {
            _cnt++;
        }
        bool _isword;
        int _cnt;
        unordered_map _nodes;
    };
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

字典树类结构设计

同样的，字典树类只是删去了映射仿函数而已

class Trie
    {
    private:
        typedef TrieNode Node;

    public:
        Trie() : _root(1, Node())
        {
        }

        void insert(const string &str)
        {
            int idx, curidx = 0;
            Node *cur = root();
            for (auto ch : str)
            {
                if (!cur->isexist(ch))
                {
                    cur = genNode(curidx);
                    cur->addnode(ch, _root.size() - 1);
                }
                curidx = cur->getNode(ch);
                cur = to_Node(curidx);
                cur->addcnt();
            }
            cur->setword();
        }
        bool search(const string &str)
        {
            int idx;
            Node *cur = root();
            for (auto ch : str)
            {
                if (!cur->isexist(ch))
                {
                    return false;
                }
                cur = to_Node(cur->getNode(ch));
            }
            return cur->isword();
        }
        Node *root()
        {
            return &_root[0];
        }
        Node *to_Node(int node_idx)
        {
            return &_root[node_idx];
        }
        Node *genNode(int curidx)
        {
            _root.emplace_back(Node());
            return &_root[curidx];
        }

    private:
        vector _root;
    };
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