Trie树

Trie树
Trie树已经在之前的文章中介绍过了。

这个题目算作是使用Trie树的又一个例子吧。为了提高效率，需要在构造过程中就进行前缀计数。

#include <iostream>
#include <cstdio>
#include <cstring> // memset
using namespace std;
// 880ms    72MB
class Trie;
class TrieNode {
public:
    TrieNode(int n = 0) : count(n) {
        memset(next, 0, sizeof(next)); // 104
    }

    friend class Trie;
private:
    int count;
    TrieNode* next[26];
};
class Trie {
public:
    Trie() { root = new TrieNode(); }
    void constructString(const char*);
    TrieNode* findString(const char*);
    int prefixStringNum(const char*);
    int prefixTrieNodeNum(const TrieNode*);
private:
    TrieNode* root;
};

int main() {
    int n, m;
    scanf("%d", &n);
    char s[11];
    Trie trie;
    while (n--) {
        scanf("%s", s);
        trie.constructString(s);
    }
    scanf("%d", &m);
    while (m--) {
        scanf("%s", s);
        cout << trie.prefixStringNum(s) << endl;
    }
    return 0;
}

void Trie::constructString(const char* s) {
    TrieNode* p = root;
    for (int i = 0; s[i]; ++i) {
        int index = s[i] - 'a';
        if (!p->next[index]) {
            p->next[index] = new TrieNode();
        }
        ++p->count;
        p = p->next[index];
    }
    ++p->count;
}
TrieNode* Trie::findString(const char* s) {
    TrieNode* p = root;
    for (int i = 0; s[i]; ++i) {
        int index = s[i] - 'a';
        if (!p->next[index]) {
            return NULL;
        }
        p = p->next[index];
    }
    return p;
}
int Trie::prefixStringNum(const char* s) {
    TrieNode* p = findString(s);
    if (!p) {
        return 0;
    }
    else {
        // return prefixTrieNodeNum(p);
        return p->count;
    }
}

int Trie::prefixTrieNodeNum(const TrieNode* p) {
    if (!p) {
        return 0;
    }
    else {
        int count = p->count;
        for (int i = 0; i < 26; ++i) {
            count += prefixTrieNodeNum(p->next[i]);
        }
        return count;
    }
}