C++でBIT（Binary Indexed Tree）を使って色付き木の部分木に含まれる異なる色の数をクエリする方法

本記事では、C++のBIT（Binary Indexed Tree、フェニック木）を活用して、色付き木の部分木に含まれる「異なる色の数」を効率的にクエリするアルゴリズムを解説します。

問題の概要

根付き木が与えられ、各ノードには配列で指定された色が割り当てられています。ここで求めたいのは、指定されたノードを根とする部分木（そのノードより下位に存在するすべてのノード）の中に、何種類の異なる色が含まれているかということです。

アルゴリズムの考え方

木構造をそのまま扱うと部分木ごとの集計が難しいため、この手法では問題を「配列上の区間クエリ」に変換して処理します。ポイントは次の3つです。

1. DFSによる木の平坦化（オイラーツアー）： DFSを実行し、各ノードの「訪問時刻（入った時刻）」と「出た時刻」を記録します。これにより、任意のノードの部分木は、訪問時刻から出発時刻までの連続する区間として表現できます。
2. 色ごとの出現位置の管理： 平坦化した配列上で、各色が出現する位置を色ごとのテーブルに保存します。
3. BITによる異なり数の管理： 各色について「現在有効な出現位置」だけをBITに登録します。クエリを訪問時刻順にソートしてオフライン処理することで、区間 [visTime[v], endTime[v]] に含まれる異なる色の数を、BITの区間和として取得できます。

C++での実装例

以下に完全な実装コードを示します。

#include<bits/stdc++.h>
#define MAXIMUM_COLOUR 1000005
#define MAXIMUM_NUMBER 100005
using namespace std;
vector<int> tree[MAXIMUM_NUMBER];
vector<int> table[MAXIMUM_COLOUR];
int isTraversing[MAXIMUM_COLOUR];
int bit[MAXIMUM_NUMBER], getVisTime[MAXIMUM_NUMBER],
getEndTime[MAXIMUM_NUMBER];
int getFlatTree[2 * MAXIMUM_NUMBER];
bool vis[MAXIMUM_NUMBER];
int tim = 0;
vector< pair< pair<int, int>, int> > queries;
// 各クエリの結果を保存する配列
int ans[MAXIMUM_NUMBER];
void update(int idx, int val) {
    while ( idx < MAXIMUM_NUMBER ) {
        bit[idx] += val;
        idx += idx & -idx;
    }
}
int queryingTree(int idx) {
    int result = 0;
    while ( idx > 0 ) {
        result += bit[idx];
        idx -= idx & -idx;
    }
    return result;
}
void preformingDFS(int v, int color[]) {
    // ノードを訪問済みとしてマーク
    vis[v] = 1;
    getVisTime[v] = ++tim;
    getFlatTree[tim] = color[v];
    vector<int>::iterator it;
    for (it=tree[v].begin(); it!=tree[v].end(); it++)
        if (!vis[*it])
            preformingDFS(*it, color);
    getEndTime[v] = ++tim;
    getFlatTree[tim] = color[v];
}
// 木に辺を追加する
void addingNewEdge(int u, int v) {
    tree[u].push_back(v);
    tree[v].push_back(u);
}
void markingFirstFind(int n) {
    for (int i = 1 ; i <= 2 * n ; i++) {
        table[getFlatTree[i]].push_back(i);
        if (table[getFlatTree[i]].size() == 1) {
            update(i, 1);
            isTraversing[getFlatTree[i]]++;
        }
    }
}
void calcQuery() {
    int j = 1;
    for (int i=0; i<queries.size(); i++) {
        for ( ; j < queries[i].first.first ; j++ ) {
            int elem = getFlatTree[j];
            update( table[elem][isTraversing[elem] - 1], -1);
            if ( isTraversing[elem] < table[elem].size() ){
                update(table[elem][ isTraversing[elem] ], 1);
                isTraversing[elem]++;
            }
        }
        ans[queries[i].second] = queryingTree(queries[i].first.second);
    }
}
// 異なる色のノード数をカウントする
void calcAllColours(int color[], int n, int qVer[], int qn) {
    preformingDFS(1, color);
    for (int i=0; i<qn; i++)
        queries.push_back(make_pair(make_pair(getVisTime[qVer[i]] , getEndTime[qVer[i]]), i) );
    sort(queries.begin(), queries.end());
    markingFirstFind(n);
    calcQuery();
    for (int i=0; i<queries.size() ; i++) {
        cout << "All distinct colours in the given tree: " << ans[i] << endl;
    }
}
int main() {
    int number = 6;
    int color[] = {0, 2, 3, 3, 4, 1};
    addingNewEdge(1, 2);
    addingNewEdge(1, 3);
    addingNewEdge(2, 4);
    int queryVertices[] = {3, 2};
    int qn = sizeof(queryVertices)/sizeof(queryVertices[0]);
    calcAllColours(color, number, queryVertices, qn);
    return 0;
}

実行結果

All distinct colours in the given tree: 1
All distinct colours in the given tree: 2

サンプルの解説

このサンプルでは、6つのノードからなる木を構築しています。ノード1〜5の色はそれぞれ {2, 3, 3, 4, 1} であり、辺は (1,2)、(1,3)、(2,4) で接続されています。

• ノード3へのクエリ： ノード3の部分木はノード3のみ（色3）なので、異なる色は1種類です。
• ノード2へのクエリ： ノード2の部分木はノード2（色3）とノード4（色4）で構成されるため、異なる色は2種類です。

計算量

DFSによる木の平坦化にO(N)、クエリのソートにO(Q log Q)、各クエリの処理はBITの更新・参照によりO(log N)ずつかかります。したがって、全体の計算量は O((N + Q) log N) となり、ノード数やクエリ数が多い場合でも効率的に処理できるのがこの手法の大きな利点です。