トライ木(Trie)は、辞書木や接頭辞木とも呼ばれるデータ構造で、特定の文字列が存在するかどうかの検索や、特定の接頭辞を持つ文字列の数を効率的に数えるために使用されます。
トライ木は接頭辞の概念を利用しており、各文字列を一文字ずつ分解して木構造のノードに格納します。例えば、"cup", "apple", "cake", "app", "blog" という単語がある場合、以下のような木構造を構築できます:
しかし、実際の実装では文字を辺ではなく、深さの大きいノードに格納するのが一般的です:
次に、完全な文字列であるかどうかを判定する方法を考えます。例えば、"ca" が完全な文字列かどうかをどう判断すればよいでしょうか?最後の文字が葉ノードかどうかで判断するのは適切ではありません。上記の例では、"app" の最後の文字は葉ノードではありません。この問題を解決するために、各ノードに終端マークを追加します。
トライ木の基本概念を理解したところで、挿入と検索の操作を実装しましょう。二次元配列 node[i][j] を使用して、ノード i の文字 j に対応する子ノードの番号を表現します。文字を直接配列の添字として使用できないため、文字を数値に変換します。例えば、a~z の文字のみが使用される場合、c - 'a' のように変換できます。
挿入と検索操作では、文字列の各文字を順に処理し、現在のノード番号を追跡する変数 current を使用します。i 番目の文字に対して、対応する子ノードが存在するか(node[current][s[i]] が 0 でないか)を確認します。存在しない場合、挿入操作では新しいノードを作成し、検索操作では false を返します。
実装コードは以下の通りです:
void addWord(string str) {
int length = str.length();
int current = 0; // 現在のノード位置
for (int i = 0; i < length; i++) {
int charIndex = str[i] - 'a';
if (node[current][charIndex] == 0) {
node[current][charIndex] = ++nodeCount; // 新しいノードを作成
}
current = node[current][charIndex];
}
isEnd[current] = true; // 単語の終わりをマーク
return;
}
bool searchWord(string str) {
int length = str.length();
int current = 0;
for (int i = 0; i < length; i++) {
int charIndex = str[i] - 'a';
if (node[current][charIndex] == 0) {
return false;
}
current = node[current][charIndex];
}
return isEnd[current];
}
次に、特定の文字列を接頭辞として持つ単語の数を数える方法を考えます。prefixCount[i] を使用して、根ノードからノード i へのパスに対応する接頭辞を持つ単語の数を記録します。挿入操作では、通過する各ノードの prefixCount をインクリメントし、検索操作では最終ノードの prefixCount を返します。
接頭辞カウントの実装は以下の通りです:
int countPrefix(string str) {
int length = str.length();
int current = 0;
for (int i = 0; i < length; i++) {
int charIndex = str[i] - 'a';
if (node[current][charIndex] == 0) {
return 0;
}
current = node[current][charIndex];
}
return prefixCount[current];
}
完全な実装コードは以下の通りです:
#include <iostream>
#include <string>
using namespace std;
const int MAX_NODES = 1000005;
const int ALPHABET_SIZE = 26;
int nodeCount;
int node[MAX_NODES][ALPHABET_SIZE];
int prefixCount[MAX_NODES];
bool isEnd[MAX_NODES];
void addWord(string str) {
int length = str.length();
int current = 0;
for (int i = 0; i < length; i++) {
int charIndex = str[i] - 'a';
if (node[current][charIndex] == 0) {
node[current][charIndex] = ++nodeCount;
}
current = node[current][charIndex];
prefixCount[current]++;
}
isEnd[current] = true;
}
bool searchWord(string str) {
int length = str.length();
int current = 0;
for (int i = 0; i < length; i++) {
int charIndex = str[i] - 'a';
if (node[current][charIndex] == 0) {
return false;
}
current = node[current][charIndex];
}
return isEnd[current];
}
int countPrefix(string str) {
int length = str.length();
int current = 0;
for (int i = 0; i < length; i++) {
int charIndex = str[i] - 'a';
if (node[current][charIndex] == 0) {
return 0;
}
current = node[current][charIndex];
}
return prefixCount[current];
}
int main() {
int wordCount, queryCount;
cin >> wordCount >> queryCount;
for (int i = 0; i < wordCount; i++) {
string str;
cin >> str;
addWord(str);
}
for (int i = 0; i < queryCount; i++) {
int queryType;
cin >> queryType;
string str;
cin >> str;
if (queryType == 1) {
cout << (searchWord(str) ? "Yes" : "No") << endl;
} else {
cout << countPrefix(str) << endl;
}
}
return 0;
}
この実装は、文字列の検索と接頭辞カウントの両方を効率的に行うことができます。トライ木の時間計算量は、操作対象の文字列の長さに比例し、検索、挿入、接頭辞カウントのいずれも O(L) で実行できます。ここで L は文字列の長さです。