P2922 [USACO08DEC]Secret Message G

P2922 [USACO08DEC]Secret Message G

给定 $n$ 个单词。对于每个询问，求单词中有多少个单词要么是询问的前缀要么询问是其前缀。

建 Trie ，对于前者，在 Trie 中找查询，路径上的单词数量即为答案；对于后者，需要知道当前子树中还有多少个单词，为了保证时间复杂度，可以记忆化搜索。

查看代码

#include <cstdio>
#include <vector>
#include <iostream>
using std::cin;
using std::vector;
const int N = 5e5 + 10;
int m, n, f[N];
int idx, cnt[N], tr[N][2];
int dfs(int x)
{
    if (f[x])
        return f[x];
    int res = cnt[x];
    for (int i = 0; i < 2; i++)
        if (tr[x][i])
            res += dfs(tr[x][i]);
    return f[x] = res;
}
void insert(vector<bool> str)
{
    int p = 0;
    for (int i = 0; i < str.size(); i++)
    {
        int t = str[i];
        if (!tr[p][t])
            tr[p][t] = ++idx;
        p = tr[p][t];
    }
    cnt[p]++;
}
int query(vector<bool> str)
{
    int res = 0;
    int p = 0;
    for (int i = 0; i < str.size(); i++)
    {
        res += cnt[p];
        int t = str[i];
        if (!tr[p][t])
            return res;
        p = tr[p][t];
    }
    return res + dfs(p);
}
int main()
{
    scanf("%d%d", &m, &n);
    for (int i = 1, a, b; i <= m; i++)
    {
        vector<bool> str;
        scanf("%d", &a);
        while (a--)
        {
            scanf("%d", &b);
            str.push_back(b);
        }
        insert(str);
    }
    for (int i = 1, a, b; i <= n; i++)
    {
        vector<bool> str;
        scanf("%d", &a);
        while (a--)
        {
            scanf("%d", &b);
            str.push_back(b);
        }
        printf("%d\n", query(str));
    }
    return 0;
}