F - Double Sum 2 Editorial
by
sounansya
まず、 \(d_k\) を以下の \(2\) つの条件
- \(1\le i\le j\le N\)
- \(A_i+A_j\) が \(2^k\) の倍数
を満たす \((i,j)\) に対する \(A_i+A_j\) の総和と定義します。すると、以下の \(2\) つの条件
- \(1\le i\le j\le N\)
- \(A_i+A_j\) が \(2\) でちょうど \(k\) 回割り切れる
を満たす \((i,j)\) に対する \(A_i+A_j\) の総和は \(d_k-d_{k+1}\) と表されるため、求める答えは \(\displaystyle \sum_{k\geq 0}\frac{d_k-d_{k+1}}{2^k}\) となります。 \(A_i+A_j\le 2\times 10^7<2^{25}\) より \(k\) の範囲は \(0\le k\le 25\) として良いです。以降は各 \(k\) に対して \(d_k\) を求めることを考えます。
\(A_i+A_j\) が \(2^k\) の倍数という条件は \(A_j\equiv -A_i\mod 2^k\) と言い換えられます。よって、 \(j\) を固定して
- \(1\le i\le j\)
- \(A_j\equiv -A_i\mod 2^k\)
を満たす \(A_i\) の個数と総和をそれぞれ \(C_j,S_j\) とすると、 \(\displaystyle d_k=\sum_{j=1}^N(C_jA_j+S_j)\) が成立します。この \(C_j\) と \(S_j\) は鍵として \(k=\left(-A_i\ \text{mod}\ 2^k\right)\) の値を、値として鍵が \(k\) となる \(A_i\) の個数と総和を持つ辞書 (map) を順番に作っていくことで高速に求めることができます。
以上を適切に実装することでこの問題を解くことができます。計算量は \(O(N\log \max A)\) です。
from collections import defaultdict
import sys
input = sys.stdin.readline
n = int(input())
a = list(map(int, input().split()))
K = 25
cur = [0] * (K + 1)
for k in range(K):
kk = 1 << k
s = defaultdict(lambda: (0, 0))
for i in a:
c1, c2 = s[-i % kk]
s[-i % kk] = c1 + 1, c2 + i
c1, c2 = s[i % kk]
cur[k] += c2 + i * c1
ans = 0
for i in range(K):
assert cur[i] % (1 << i) == 0
ans += (cur[i] - cur[i + 1]) >> i
print(ans)
posted:
last update: