C++で文字のストリームをチェックする：トライ木によるStreamCheckerの実装

文字が順次入力されるストリームの中で、直近の文字列が指定した単語リストに含まれるかどうかを即座に判定したい場面はよくあります。本記事では、C++でトライ木（Trie）を用いてこのような StreamChecker クラスを実装する方法を解説します。

問題の概要

次のような StreamChecker クラスを実装することを考えます。

• StreamChecker(words) − コンストラクタです。与えられた単語リストでデータ構造を初期化します。

• query(letter) − ある k ≥ 1 に対して、これまでに照会された最後の k 文字（今回照会した文字を含み、古いものから新しいものへと並べたもの）が、単語リスト内のいずれかの単語と一致するときに true を返します。

たとえば、入力の単語リストが ["ce", "g", "lm"] であるとします。このとき、[a, b, c, e, f, g, h, i, j, k, l, m] の順に query を何度も呼び出すと、e、g、m を渡したときだけ true となり、それ以外では false が返ります。

アルゴリズム：トライ木による解法

この問題は、単語リストをトライ木として構築し、マッチングの途中経過を管理することで効率的に解けます。手順は以下のとおりです。

• ノード構造体 Node を定義します。26 個の子ノードポインタ配列と、isEnd フラグを持ちます。

• 初期状態では isEnd は false、子配列はすべて NULL で埋められています。

• ルートノード head を用意します。

• マッチング候補を保持するためのノード配列 waitList を作成します。

• 関数 insertNode() を定義します。引数として head と文字列 s を受け取ります。

  • curr := head とします。

  • i := 0 から s のサイズまで繰り返します。

    • x := s[i] とします。

    • curr の child[x - 'a'] が NULL の場合は、新しい Node を生成して代入します。

    • curr := curr->child[x - 'a'] として子ノードへ進みます。

  • 最後に、curr の isEnd を true に設定します。

• コンストラクタでは以下を行います。

  • head := 新しい Node を生成します。

  • words の各要素について insertNode(head, words[i]) を呼び出してトライ木を構築します。

• 関数 query() を定義します。引数として文字 x を受け取ります。

  • ノード配列 temp を用意します。

  • head の child[x - 'a'] が存在する場合、head を waitList の末尾に追加します（ここが新しいマッチングの開始点になります）。

  • ret := false で初期化します。

  • waitList の各要素 i について処理します。

    • curr := waitList[i] とします。

    • curr の child[x - 'a'] が存在する場合は、curr をその子ノードへ進め、curr を temp の末尾に追加し、ret := ret OR curr->isEnd とします。

  • swap(temp, waitList) で待機リストを更新します（進められなかった候補は自動的に脱落します）。

  • ret を返します。

仕組みのポイント

waitList には、「現在の文字位置から始まる可能性のある接尾辞のマッチング候補」が保持されています。毎回の query で各候補を 1 文字だけ進め、トライ木の終端ノード（isEnd == true）に到達できれば、過去のどこかの時点から始まる文字列が単語リストに一致したことになります。文字履歴をすべて保存して毎回全文比較する必要がないため、非常に効率的な手法です。

実装例

#include <bits/stdc++.h>
using namespace std;
struct Node {
    Node* child[26];
    bool isEnd;
    Node(){
        isEnd = false;
        for (int i = 0; i < 26; i++)
        child[i] = NULL;
    }
};
class StreamChecker {
    public:
    Node* head;
    vector<Node*> waitList;
    void insertNode(Node* head, string& s){
        Node* curr = head;
        for (int i = 0; i < s.size(); i++) {
            char x = s[i];
            if (!curr->child[x - 'a']) {
                curr->child[x - 'a'] = new Node();
            }
            curr = curr->child[x - 'a'];
        }
        curr->isEnd = true;
    }
    StreamChecker(vector<string>& words){
        head = new Node();
        for (int i = 0; i < words.size(); i++) {
            insertNode(head, words[i]);
        }
        Node* curr = head;
    }
    bool query(char x){
        vector<Node*> temp;
        if (head->child[x - 'a']) {
            waitList.push_back(head);
        }
        bool ret = false;
        for (int i = 0; i < waitList.size(); i++) {
            Node* curr = waitList[i];
            if (curr->child[x - 'a']) {
                curr = curr->child[x - 'a'];
                temp.push_back(curr);
                ret |= curr->isEnd;
            }
        }
        swap(temp, waitList);
        return ret;
    }
};
main(){
    vector<string> v = {"ce","g","lm"};
    StreamChecker ob(v);
    cout << (ob.query('a')) << endl;
    cout << (ob.query('b')) << endl;
    cout << (ob.query('c')) << endl;
    cout << (ob.query('e')) << endl;
    cout << (ob.query('f')) << endl;
    cout << (ob.query('g')) << endl;
    cout << (ob.query('h')) << endl;
    cout << (ob.query('i')) << endl;
    cout << (ob.query('j')) << endl;
    cout << (ob.query('k')) << endl;
    cout << (ob.query('l')) << endl;
    cout << (ob.query('m'));
}

入力

"ce","g","lm", query(),

出力

0 0 0 1 0 1 0 0 0 0 0 1