Acwing 835 Trie字符串统计

#include<algorithm>
#include<iostream>
#include<string>
using namespace std;
const int N=1e5+10;
int son[N][26],cnt[N];
int idx;
char str[N];
void insert()
{
    int p=0;//根节点
    for(int i=0;str[i];i++)
    {
        int u=str[i]-'a';//把字母转化成数字
        if(!son[p][u]) son[p][u]=++idx;//增加根p的儿子节点
        p=son[p][u];//改变根节点
    }
    cnt[p]++;//记录字符串个数
}
int query()
{
    int p=0;//根节点
    for(int i=0;str[i];i++)
    {
        int u=str[i]-'a';
        if(!son[p][u]) return 0;//无儿子节点，返回0
        p=son[p][u];//更新根节点p
    }
    return cnt[p];//返回字符串个数
}
int main()
{
    int m;
    cin>>m;
    while(m--)
    {
        char op[2];
        cin>>op>>str;
        if(op[0]=='I')
        insert();
        else
        {
            cout<<query()<<endl;
        }
    }
    return 0;
}