yukicoder : No.641 Team Contest Estimation

制限時間があると焦る。
https://yukicoder.me/problems/no/641

問題

N要素の整数列A[i]と、整数Kが与えられる。
0≦x＜2^Kであるxに対し、f(x) = (A[0]^x) + (A[1]^x) + ～ + (A[N-1]^x) とおく。
f(0)～f(2^K-1)の平均をμ、標準偏差をσと置くとき、μ*2^Kとσ^2*4^Kを答えよ。

解法

A[0]^xの部分は、A[0]の値に関わらずxを変化させると0～(2^K-1)が1回ずつ登場する。
よってμ*2^Kはf(x)の総和であり、その値はN*(2^K)*(2^K-1)/2なので容易に定まる。

標準偏差はσ^2*4^K = (2^K)*sum(f(x)^2) - (μ*2^K)^2なので、sum(f(x)^2)を求めることを考える。
定番テクとして、数列の二乗和を考えるとき、一乗和と二乗和を順次更新していくテクが使える。

(A[i]^x)のあるビットが、xに対しどう変わるかを考える。
A[i]のうち(2^y)のビットが立っているのがa個、立っていないのがN-a個とする。
xを総当たりすると、f(x)において2^(K-1)通りのパターンで(2^y)のビットがa個たち、残り2^(K-1)通りのパターンで(2^y)のビットが(N-a)個立つ。
この情報を各ビット毎に一乗和・二乗和ともに足しこんでいけばよい。

int N,K;
ll A[101010];
ll mo=1000000009;
ll ave;
int cnt[61];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>A[i];
		FOR(j,K) cnt[j] += (A[i]>>j)&1;
	}
	
	if(K==0) return _P("0\n0\n");
	
	ave=((1LL<<K)-1)%mo*((1LL<<(K-1))%mo)%mo;
	ave = ave*N%mo;
	
	cout<<ave<<endl;
	
	ll S1=0,S2=0;
	FOR(i,K) {
		ll a=cnt[i]*((1LL<<i)%mo)%mo;
		ll b=(N-cnt[i])*((1LL<<i)%mo)%mo;
		(S2 += (S1*(a+b)+(a*a+b*b)%mo*((1LL<<(K-1))%mo)%mo))%=mo;
		(S1+=(a+b)*((1LL<<(K-1))%mo))%=mo;
	}
	
	S2 = S2*((1LL<<K)%mo)%mo - ave*ave%mo;
	cout<<(S2%mo+mo)%mo<<endl;
	
	
}

まとめ

二乗和を求める問題はいつも混乱する。