hihocoder1014 Trie树

问题描述:

Trie树在字符串处理中的应用十分重要，最典型的应用就是输入法和搜索引擎中的字符串自动补全功能。其核心思想是用一颗树来存储一个字典，树的每一条边表示单词的一个字符，在每个节点上记录以从根节点到当前节点所经过的路径为前缀的字符串个数。

利用字典树，可以实现O(log(n))的单词插入、单词查询、查询以某个前缀开头的字符串数目等。

题目链接：http://hihocoder.com/problemset/problem/1014

我的代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

#define MAX_CH 26
#define MAX_NODE 1000005

int nodecnt;

struct TrieNode
{
    int cnt;
    TrieNode *p[MAX_CH];
    void init()
    {
        cnt = 0;
        for(int i=0; i<MAX_CH; ++i) p[i] = NULL;
    }
    int query(char *s)
    {
        int idx = s[0]-'a';
        if(p[idx]!=NULL)
        {
            if(s[1]=='') return p[idx]->cnt;
            else return p[idx]->query(s+1);
        }
        else
        {
            return 0;
        }
    }
}t_node[MAX_NODE];

struct Trie
{
    TrieNode *rt;
    void init()
    {
        rt = &t_node[0];
        rt->init();
        nodecnt = 1;
    }
    void insert(char *s)
    {    
        TrieNode *now = rt;    
        for(int i=0; s[i]; ++i)
        {
            int idx = s[i]-'a';
            if(now->p[idx]==NULL)
            {
                now->p[idx] = &t_node[nodecnt++];
                now->p[idx]->init();
            }
            ++(now->p[idx]->cnt);
            now = now->p[idx];
        }
    }
    int query(char *s)
    {
        return rt->query(s);
    }
}trie;

int main()
{
    int n, m;
    char s[20];
    while(scanf("%d", &n)!=EOF)
    {
        trie.init();
            
        while(n--)
        {
            scanf("%s", s);
            trie.insert(s);
        }
        
        scanf("%d", &m);
        while(m--)
        {
            scanf("%s", s);
            printf("%d
", trie.query(s));
        }
    }    
    return 0;
}