AcWing 835. Trie字符串统计（Trie）

#include <bits/stdc++.h>
using namespace std;

const int N = 300000;
int son[N][26], cnt[N], idx;
int n;
char str[N];

void insert(char *str) {
    int p = 0;
    for(int i = 0; str[i]; i++) {
        char c = str[i];
        if(!son[p][c-'a']) son[p][c-'a'] = ++idx;
        p = son[p][c-'a'];
    }
    cnt[p]++;
}

int query(char *str) {
    int p = 0;
    for(int i = 0; i < str[i]; i++) {
        char c = str[i];
        if(!son[p][c-'a']) return 0;
        p = son[p][c-'a'];
    }
    return cnt[p];
}

int main() {
    scanf("%d",&n);
    while(n -- ) {
        char op[2];
        scanf("%s%s",op,str);
        if(op[0] == 'I') insert(str);
        else cout << query(str) << endl;;
    }
    return 0;
}