CodeChef February Challenge 2019: Xor Decomposition - 日々ｄｒｄｒする人のメモ

CodeChef February Challenge 2019 の問題: Xor Decomposition (Code: XDCOMP)
問題ページ: https://www.codechef.com/FEB19A/problems/XDCOMP

問題概要

1から ${N}$ の番号が付いた ${N}$ 頂点の木と一つの整数 ${x}$ が与えられる。
グラフの各頂点には重みが付いており、頂点 ${i}$ の重みは ${a_i}$ である。

ここから、木の辺をいくつか削除する。そして、できた森に含まれる各木について、その木に含まれる頂点すべての重みのXORをとった値が ${x}$ になるようにしたい。

このような辺の削除の仕方は何通り存在するか。 ${10^9 + 7}$ で割った余りで答えよ。

制約

${1 \le N \le 10^5}$
${0 \le x \le 10^9}$
各iについて ${0 \le a_i \le 10^9}$

解法

木dpをする。

まず、森を分割して全ての木でXORを ${x}$ にできた時のことを考えると、元々の木全体のXORの値が 0もしくは ${x}$ になってないと分割が達成できないことがわかる。
そのため、はじめに木全体のXORが0もしくは ${x}$ でない時の通り数は0になる。

また、木全体のXORの値が0の時は偶数個の木、 ${x}$ の時は奇数個の木に分割する必要があることがわかる。
そのため、今回は偶数個に分割できる通り数、もしくは奇数個に分割できる通り数を数えていくことで、解を計算する。

この奇数、偶数の通り数計算は木dpで行う。
具体的には、木を根付き木として考え、以下のdpを計算する。
${dp[v][k] =}$ 頂点v以下の部分木について、(頂点vを含む木を分割途中として個数から除くケースも含めて) XORが ${x}$ になる(k = 0: 偶数, k = 1: 奇数)個の木に分割できる個数

f:id:smijake3:20190217010543p:plain — 分割途中の例: 頂点v(赤い頂点)を含む木を分割途中とし、この木を個数に含めずに分割された木が3個(奇数個)であると考える

後はこのdpを計算した後、全体のXORが 0 か x かに応じて、 ${dp[v][0], dp[v][1]}$ のどっちかを解として出力すればよい。

気をつけるべきケースとして、 ${x = 0}$ のケースがある。
これは、分割が奇数個、偶数個のどっちにもなりうるので注意する必要がある。

計算量は ${O(N)}$

実装

提出コード (PyPy3): Solution: 22799167 | CodeChef

import sys
sys.setrecursionlimit(10**6)
readline = sys.stdin.readline
write = sys.stdout.write

MOD = 10**9 + 7

N, X = map(int, readline().split())
G = [[] for i in range(N)]
*A, = map(int, readline().split())
for i in range(N-1):
    u, v = map(int, readline().split())
    G[u-1].append(v-1)
    G[v-1].append(u-1)
Y = 0
for a in A:
    Y ^= a

# 全体のXORがX, 0以外の通り数は0
if X != Y != 0:
    write("0\n")
    exit(0)

B = [0]*N
# 頂点v以下の部分木に関するdfs
def dfs(v, p):
    # a = dp[v][0], b = dp[v][1] に該当
    a = 1; b = 0
    # s: 頂点v以下の部分木のXORの値
    s = A[v]
    for w in G[v]:
        if w == p:
            continue
        ra, rb = dfs(w, v)
        s ^= B[w]
        # 偶奇を合わせながら個数を計算
        a, b = (ra*a + rb*b) % MOD, (ra*b + rb*a) % MOD
    # ここまでは、分割後において頂点vを含む部分木のXORが x で
    #  分割の個数に含まれるケースを考えてないので、ここで計算する
    if p != -1:
        if s == X:
            # 頂点v以下の部分木全体のXORの値が x
            # (頂点v以下の部分木は、XORがxの部分木にちょうど奇数個で分割できる)
            # dp[v][0] += dp[v][1] をする
            b = (a + b) % MOD
            if X == 0:
                # x = 0 の時は特殊ケース
                # dp[v][0] = dp[v][1] = dp[v][0] + dp[v][1] をする
                a = b
        elif s == 0:
            # 頂点v以下の部分木全体のXORの値が 0
            # (頂点v以下の部分木は、XORがxの部分木にちょうど偶数個で分割できる)
            # dp[v][1] += dp[v][0] をする
            a = (a + b) % MOD
    B[v] = s
    return a, b

a, b = dfs(0, -1)
# この実装では、分割後の森において、根頂点(0)を含む部分木が x になる分を数えてないので、
# 奇数と偶数を逆に出力している
write("%d\n" % (((a+b) % MOD if X == 0 else a) if Y == X else b))