🔢

ABC365 - E 復習解説

2024/08/04に公開

競技プログラミング

AtCoder

abc365

tech

自分の知識が不足していたせいもあり公式解説を読んでも最初理解できなかったので ABC365 E を復習して自分なりにまとめてみました。初学者の方の参考になれば幸いです。

E - Xor Sigma Problem

XOR に関する問題は各桁を独立して考えることができることが典型ですので今回も独立して考えてみます。

最初に2進数にした時の桁 $p$ の $A$ の $i$ 番目の要素を $B_{i}$ とおきます。この時、2進数なので $B_{i}$ は $0$ または $1$ となります。

例えば $p = 3$ の場合、以下のように $B$ を計算できます。

p = 3
B = [(a >> p) & 1 for a in A]

次に、累積和と同様の考え方で $B_{1}, B_{2}, ..., B_{N}$ の要素を前の方から XOR を取っていく累積 XOR を求めます。初項として $0$ も追加しておきます。(ちなみに累積 XOR 自体は桁ごとに考えずに $A_{i}$ の値でやっても成立します)

from itertools import accumulate

acc = [0] + list(accumulate(B, lambda x, y: x ^ y))

さて、 E 問題で求めようとしている式の値は累積 XOR acc を用いるとすべての区間 $(l, r)$ を足せばよく、以下のように表すことができます。同じ値をもう一度 XOR とることで累積和で言う引き算のようにできるのがポイントです。

ans = 0
for l in range(N - 1):
    for r in range(l, N):
        ans += acc[r] ^ acc[l - 1]

しかし、このままでは愚直に計算すると累積 XOR を用いたところで計算量は $O(N^{2})$ となってしまいます。

ここで、累積 XOR の値は $0$ または $1$ のいずれかしか取らないことを再度考えてみます。

そう考えるとある区間 $(l, r)$ の acc[r] ^ acc[l - 1] は

0 ^ 0 = 0
0 ^ 1 = 1
1 ^ 0 = 1
1 ^ 1 = 0

となり、 $(l, r)$ は異なる値を選んだときにだけ $1$ となると言えます。

💡 この4パターンしかない状態に持っていくために各桁で考えたほうがいいわけですね。おもしろいですね。

したがって、累積 XOR を取りながら累積 XOR の中の $0$ と $1$ の数をそれぞれカウントしておき(最後にカウントしても OK です)、 $0$ と $1$ のそれぞれの数をかけ算したものが桁 $p$ が 1 となる区間 $(l, r)$ の組み合わせと言えます。

2進数の桁 $p$ の値は10進数だと $1 << p$ と言えるので求める答えに $(組み合わせの数) \times 1 << p$ を加えましょう。for 文が2つとも消えてしまいました。

count_0 = ... # 累積 XOR の 0 の数をカウントしておく
count_1 = ... # 累積 XOR の 1 の数をカウントしておく
ans = count_0 * count_1 * (1 << p)

💡 私の中ではこういう考えですべての組み合わせを計算しなくてよくするという発想がなかったのですごい勉強になりました。

最後に

ans = 0
for l in range(N - 1):
    for r in range(l, N):
        ans += acc[r] ^ acc[l - 1]

の式は l = r = 1 の時に acc[1] ^ acc[0] となりこれは本来ないはずの要素1つだけを区間して余計に足していることになるので答えから sum(A) を引いておきましょう。

ans -= sum(A)

なお $1 \le A_{i} \le 10^{8}$ という制約から $p$ は 30 桁まで考えたら十分です。
桁ごとに累積 XOR を計算するところが一番重いので計算量は $O(30\timesN)$ です。

以上を Python で実装すると以下のようになります。
(累積 XOR acc は考え方としては必要ですがコード自体には最終的に不要のため出てきません)

N: Final = int(input())
A: Final[list[int]] = list(map(int, input().split()))


def f(p):
    """2進数にした時のp桁目についてリストAの中の要素の0,1の数を数える"""
    counts = [1, 0] # 累積 XOR として初項に 0 を入れた状態を考えているので 0 のカウントを 1 増やしておく
    x = 0
    for a in A:
        x ^= (a >> p) & 1
        counts[x] += 1
    return counts


M = 30
ans = -sum(A)  # 先に余分に計算してしまう sum(A) を引いています
for p in range(M):
    count_0, count_1 = f(p)
    ans += count_0 * count_1 * (1 << p)
print(ans)

E - Xor Sigma Problem

Discussion