F - Double Sum 2 Editorial by shakayami

分割統治

2つの数列\(A,B\)に対して

\[F(A,B):=\sum_{i=1}^{N}\sum_{j=1}^{M} f(A_i+B_j)\]

が解ければ十分

多分分割統治みたいなことをする

数列\(A\)\(A_{odd}\cup A_{even}\)というように偶数と奇数に分ける.\(B\)についても\(B_{odd}\cup B_{even}\)と分ける.

このとき,

\[F(A,B)=F(A_{even},B_{even})+F(A_{odd},B_{even})+F(A_{even},B_{odd})+F(A_{odd},B_{odd})\]

である.

\(F(A_{odd},B_{even})\)については, 偶数と奇数の和は奇数なので,

\(f(A_{odd,i}+B_{even,j})=A_{odd,i}+B_{even,j}\)である.

\(F(A_{even},B_{odd})\)も同様

\(F(A_{even},B_{even})\)については,

\(f(A_{even,i}+B_{even,j})=f(A_{even,i}/2+B_{even,j}/2)\)

なので,

\(F(A_{even},B_{even})=F(A_{even}/2,B_{even}/2)\)

\(F(A_{odd},B_{odd})\)については,\(f(A_{odd,i}+B_{odd,j})=f(\lfloor A_{odd,i}/2\rfloor +\lfloor B_{odd,j}/2\rfloor +1)\)なので,\(F(A_{odd},B_{odd})=F(\lfloor A_{odd}/2\rfloor ,1+\lfloor B_{odd}/2\rfloor )\)

これを再帰関数として実装すればOK

def solve(A,B):
    '''
    sum_{i,j} f(A_i+B_j)
    '''
    if len(A)==0 or len(B)==0:
        return 0
    A_even=[a for a in A if a%2==0]
    A_odd=[a for a in A if a%2==1]
    B_even=[b for b in B if b%2==0]
    B_odd=[b for b in B if b%2==1]
    AeBe=solve([a//2 for a in A_even],[b//2 for b in B_even])
    AeBo=len(B_odd)*sum(A_even)+len(A_even)*sum(B_odd)
    AoBe=len(A_odd)*sum(B_even)+len(B_even)*sum(A_odd)
    AoBo=solve([a//2 for a in A_odd],[1+b//2 for b in B_odd])
    return AeBe+AeBo+AoBe+AoBo
def f(x):
    while(x%2==0):
        x//=2
    return x
N=int(input())
A=[int(i) for i in input().split()]
ans=solve(A,A)
for i in range(N):
    ans+=f(2*A[i])
ans//=2
print(ans)

posted:
last update: