C++で木を分割した際の2つの部分木のビット単位ORが等しくなる辺の数を求める方法

概要

m個のノードからなる木があり、それぞれのノードには数値が割り当てられているとします。この木の任意の辺を切断すると、2つの新しい木が生成されます。ここでの課題は、ある辺を切断した後に構成される2つの木それぞれのノード値のビット単位OR（Bitwise OR）が等しくなるような辺の数を数えることです。なお、各ノードの値は10^6以下であることが保証されています。

入力例

values[]={1, 3, 1, 3}
    1
   / | \
  2  3  4

出力例

2

この例では、ノード1とノード2をつなぐ辺を切断すると、結果として得られる2つの木のビット単位ORはどちらも3になります。同様に、ノード1とノード4をつなぐ辺も切断可能です。したがって、条件を満たす辺は合計2本となります。

アルゴリズムの考え方

この問題は、単純なDFS（深さ優先探索）を用いて効率的に解くことができます。

ノードの値が10^6以下であるため、その値は22個のバイナリビットで表現できます。したがって、ノード値のビット単位ORも22ビット以内に収まります。

基本的なアプローチは以下の通りです。

• 各部分木について、すべてのノード値の中で「各ビットがセット（1）になっている回数」を求めます。
• 各辺について、ビット0から21までのそれぞれのビットに関して、「そのビットがセットされた数値の個数」が、切断後の2つの木において両方ともゼロであるか、両方ともゼロより大きいかを検証します。
• すべてのビットでこの条件が満たされていれば、その辺を答えとしてカウントします。

これは直感的にも妥当です。あるビットが一方の木にのみ存在する場合、そのビットは片方の木のOR結果に現れ、もう片方には現れないため、2つのORが一致しなくなるからです。

実装例

// C++による実装
#include<bits/stdc++.h>
using namespace std;
int m1[1000],x1[22];
// 各部分木の全ノード値における、
// 各ビットがセットされている回数を格納する配列
int a1[1000][22];
vector<vector<int>> g;
int ans1 = 0;
// 単純なDFSを実行する関数
void dfs(int u1, int p1){
    for (int i=0;i<g[u1].size();i++) {
        int v1 = g[u1][i];
        if (v1 != p1) {
            dfs(v1, u1);
            // ノードvを根とする部分木の全値における
            // 各ビットのセット回数を集計
            for (int i = 0; i < 22; i++)
                a1[u1][i] += a1[v1][i];
        }
    }
    // 各ビットについて、そのビットがセットされた数値の数が
    // 切断後の2つの木で両方ともゼロ、または
    // 両方ともゼロより大きいかどうかを検証
    int pp1 = 0;
    for (int i = 0; i < 22; i++) {
        if (!((a1[u1][i] > 0 && x1[i] - a1[u1][i] > 0)
            || (a1[u1][i] == 0 && x1[i] == 0))) {
            pp1 = 1;
            break;
        }
    }
    if (pp1 == 0)
        ans1++;
}
// ドライバーコード
int main(){
    // ノード数
    int n1 = 4;
    // int n1 = 5;
    // 木を格納するためのベクター
    g.resize(n1+1);
    // ノードの値を格納する配列
    m1[1] = 1;
    m1[2] = 3;
    m1[3] = 1;
    m1[4] = 3;
    /* m1[1] = 2;
    m1[2] = 3;
    m1[3] = 32;
    m1[4] = 43;
    m1[5] = 8;*/
    // 木全体の全値における各ビットの
    // セット回数を格納する配列の初期化
    for (int i = 1; i <= n1; i++) {
        int y1 = m1[i];
        int k1 = 0;
        // ノードiの値におけるセットビットを判定
        while (y1 != 0) {
            int p1 = y1 % 2;
            if (p1 == 1) {
                x1[k1]++;
                a1[i][k1]++;
            }
            y1 = y1 / 2;
            k1++;
        }
    }
    // 辺を追加
    g[1].push_back(2);
    g[2].push_back(1);
    g[1].push_back(3);
    g[3].push_back(1);
    g[1].push_back(4);
    g[4].push_back(1);
    //g[1].push_back(5);
    //g[5].push_back(1);
    dfs(1, 0);
    cout<<(ans1);
}

出力

2

計算量の観点では、DFSによって全ノードを一度訪問し、各ノードで22ビット分の集計を行うため、全体の時間計算量はO(N × 22)、つまりノード数Nに対して線形時間で処理できます。これにより、大規模な木でも効率的に条件を満たす辺を見つけることが可能です。