けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 045 C - エックスオア多橋君 (試験管青色)

XOR について a ^ a = 0 な性質を上手く使う問題として。

問題へのリンク

問題概要

 N 頂点の重み付きツリーが与えられる。

整数  X が与えられ、ツリー上のパスであってパス上の重みの XOR 和が  X となるものが何個あるかを数えよ。

制約

  •  N \le 10^{5}

考えたこと

根を 1 つ決めて根付き木にしてしまう。そして

  • sum[ v ] := 根から v までのパスの XOR 和

としておくと、なんと二頂点 a, b 間のパスの XOR 和は

  • sum[ a ] ^ sum[ b ]

と簡単に表せる。a と b の LCA より上の部分については XOR をとって相殺されるからだ。

よって、各頂点 v に対して

  • sum の値が W ^ sum[ v ] になっている頂点が (自分自身以外に) 何個あるか

を求めて足してあげて、最後に 2 で割れば OK。

#include <iostream>
#include <vector>
#include <map>
using namespace std;
using Edge = pair<int,int>;

int N, X;
vector<vector<Edge> > tree;
vector<int> sum;

void dfs(int v, int len, int p = -1) {
    sum[v] = len;
    for (auto e : tree[v]) {
        if (e.first == p) continue;
        dfs(e.first, len ^ e.second, v);
    }
}

int main() {
    cin >> N >> X;
    tree.assign(N, vector<Edge>());
    for (int i = 0; i < N-1; ++i) {
        int a, b, c; cin >> a >> b >> c; --a; --b;
        tree[a].push_back( make_pair(b, c) );
        tree[b].push_back( make_pair(a, c) );
    }
    sum.assign(N, -1);
    dfs(0, 0);
        
    long long res = 0;
    map<int, long long> con;
    for (int i = 0; i < N; ++i) con[sum[i]]++;
    for (int v = 0; v < N; ++v) {
        long long num = con[sum[v] ^ X];
        if (X != 0) res += num;
        else res += num-1;
    }
    cout << res/2 << endl;
}