Trie树

Posted on | In |

Trie树模板

题目来源：http://hihocoder.com/problemset/problem/1014 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
#include <iostream>
#include <string>
using namespace std;
enum NODE_TYPE {
    COMPLETED,
    UNCOMPLETED
};
class TrieNode {
public:
    TrieNode* childNodes[26];
    int freq;
    char nodeChar;
    enum NODE_TYPE type;
    TrieNode(char ch) {
        this->nodeChar = ch;
        this->freq = 0;
        for(int i = 0; i < 26; i++) {
            this->childNodes[i] = NULL;
        }
        this->type = UNCOMPLETED;
    }
    void add_freq();
    void set_type(enum NODE_TYPE);
};
class Trie {
public:
    TrieNode* root;
    Trie() {
        root = new TrieNode(' ');
    }
    void insert(string str);
    bool find(string str);
    int query(string str);
};
int charToIndex(char ch) {
    return ch - 'a';
}
void Trie::insert(string str) {
    TrieNode *current_node = this->root;
    for(int i = 0; i < str.length(); i++) {
        int index = charToIndex(str[i]);
        char ch = str[i];
        if(current_node->childNodes[index] == NULL) {
            current_node->childNodes[index] = new TrieNode(ch);
        }
        current_node->childNodes[index]->add_freq();
        if(i == str.length()-1) {
            current_node->childNodes[index]->set_type(COMPLETED);
        } else if(current_node->childNodes[index]->type != COMPLETED){
            current_node->childNodes[index]->set_type(UNCOMPLETED);
        }
        current_node = current_node->childNodes[index];
    }
}
bool Trie::find(string str) {
    TrieNode *current_node = this->root;
    int i = 0;
    while(i < str.length()) {
        char ch = str[i];
        int index = charToIndex(ch);
        if(current_node->childNodes[index] == NULL) {
            break;
        }
        i++;
        current_node = current_node->childNodes[index];
    }
    cout << current_node->nodeChar << " " << current_node->type << endl;
    return current_node->type == COMPLETED && i == str.length();
}
int Trie::query(string str) {
    TrieNode *current_node = this->root;
    int i = 0;
    int count;
    while(i < str.length()) {
        char ch = str[i];
        int index = charToIndex(ch);
        if(current_node->childNodes[index] == NULL) {
            break;
        }
        i++;
        current_node = current_node->childNodes[index];
    }
    if(i == str.length()){
        return current_node->freq;
    }
    return 0;
}
void TrieNode::set_type(enum NODE_TYPE t) {
    this->type = t;
}
void TrieNode::add_freq() {
    this->freq++;
}
int main() {
    Trie trie;
    int loop;
    cin >> loop;
    while(loop > 0) {
        string str;
        cin >> str;
        trie.insert(str);
        loop--;
    }
    cin >> loop;
    while(loop > 0) {
        string str;
        cin >> str;
        cout << trie.query(str) << endl;
        loop--;
    }
    return 0;
}