字典树的实现

字典树的实现

题目描述

字典树又称为前缀树或者Trie树，是处理字符串常用的数据结构。假设组成所有单词的字符仅是‘a’～‘z’，请实现字典树的结构，并包含以下四个主要的功能。void insert(String word)：添加word，可重复添加；void delete(String word)：删除word，如果word添加过多次，仅删除一次；boolean search(String word)：查询word是否在字典树中出现过(完整的出现过，前缀式不算)；int prefixNumber(String pre)：返回以字符串pre作为前缀的单词数量。现在给定一个m，表示有m次操作，每次操作都为以上四种操作之一。每次操作会给定一个整数op和一个字符串word，op代表一个操作码，如果op为1，则代表添加word，op为2则代表删除word，op为3则代表查询word是否在字典树中，op为4代表返回以word为前缀的单词数量（数据保证不会删除不存在的word）。

输入描述：

输入包含多行，第一行一个整数m$(1\le m\le 10^5)$，代表操作次数。接下来m行，每行包含一个整数op$(1\le op\le 4)$和一个字符串word$(1\le lengt{h}_{word}\le 20)$。

输出描述：

对于每次操作，如果op为3时，如果word在字典树中，请输出“YES”，否则输出“NO”；如果op为4时，请输出返回以word为前缀的单词数量，其它情况不输出。

示例1

输入

7
1 qwer
1 qwe
3 qwer
4 q
2 qwer
3 qwer
4 q

输出

YES
2
NO
1

---

题解：

字典树的模板题。关于字典树的定义不做介绍。本题我们需要记录有多少个单词共用某个节点，以及多少个单词以某个节点结尾。我习惯性使用二维数组表示 Trie 树，就是开辟的空间较大。

代码：

#include <cstdio>

using namespace std;

const int N = 2000010;

int trie[N][26];
int path[N];
int num[N];
int idx;

void insert( char *s ) {
    int p = 0;
    for ( int i = 0; s[i]; ++i ) {
        int u = s[i] - 'a';
        int &t = trie[p][u];
        if ( !t ) t = ++idx;
        ++path[t];
        p = t;
    }
    ++num[p];
}

bool search( char *s ) {
    int p = 0;
    for ( int i = 0; s[i]; ++i ) {
        int u = s[i] - 'a';
        int &t = trie[p][u];
        if ( !t ) return false;
        p = t;
    }
    return num[p] != 0;
}

void _delete( char *s ) {
    if ( !search( s ) ) return;
    int p = 0;
    for ( int i = 0; s[i]; ++i ) {
        int u = s[i] - 'a';
        int &t = trie[p][u];
        if ( path[t]-- == 1 ) {
            t = 0;
            return;
        }
        p = t;
    }
    --num[p];
}

int prefixNumber( char *s ) {
    int p = 0;
    for ( int i = 0; s[i]; ++i ) {
        int u = s[i] - 'a';
        int &t = trie[p][u];
        if ( !t ) return 0;
        p = t;
    }
    return path[p];
}

int main( void ) {
    int n, op;
    char s[21];
    scanf("%d", &n);
    while ( n-- ) {
        scanf("%d %s", &op, s);
        if ( op == 1 ) insert( s );
        else if ( op == 2 ) _delete( s );
        else if ( op == 3) puts( search( s ) ? "YES" : "NO" );
        else printf("%d\n", prefixNumber( s ) );
    }
    return 0;
}