AtCoder ARC #136 : D - Without Carry

Cで手間取りすぎた。
https://atcoder.jp/contests/arc136/tasks/arc136_d

問題

N個の100万未満の非負整数からなる数列Aが与えられる。
数列の要素の対のうち、足したとき繰り上がりが生じないものは何通りか。

解法

f(n) := Aのうち、各桁において、d桁目がdのi桁目以下という条件をすべて満たすものの数
とする。
これが求められれば、A[i]に対し、f(999999-A[i])個の要素を対とすることができるので、iを総当たりすればよくなる。
(正確には、(A[i],A[i])というペアがカウントされる分を引くのと、同じペアが二重カウントされるので最後に2で割る必要がある)

あとはf(n)を求めよう。
2つの方法が考えられる。

• nは6桁なので、1桁ごとに分解して10*10*10*10*10*10の6次元配列を考え、包除原理の要領で累積和を取る。
• 高速ゼータ変換の要領で累積和を取る

以下のコードは前者。
前者の方がO(2^|n|)程度重いが一応間に合う。

int N;
int A[1010101],B[1000000][6];
int C[1010101];
int S[10][10][10][10][10][10];

void solve() {
	int i,j,k,l,x,y; string s;
	
	cin>>N;
	ll ret=0;
	FOR(i,N) {
		cin>>A[i];
		x=A[i];
		C[A[i]]++;
		FOR(j,6) {
			B[i][j]=x%10;
			x/=10;
		}
		S[B[i][0]][B[i][1]][B[i][2]][B[i][3]][B[i][4]][B[i][5]]++;
	}
	int r[6],t[6];
	FOR(r[0],10)
	FOR(r[1],10)
	FOR(r[2],10)
	FOR(r[3],10)
	FOR(r[4],10)
	FOR(r[5],10) {
		int sum=0;
		int mask;
		FOR(mask,1<<6) if(mask) {
			int m=-1;
			int ng=0;
			FOR(i,6)  {
				t[i]=r[i];
				if(mask&(1<<i)) {
					t[i]--;
					m=-m;
					if(t[i]<0) ng=1;
				}
			}
			if(ng) continue;
			sum+=m*S[t[0]][t[1]][t[2]][t[3]][t[4]][t[5]];
		}
		S[r[0]][r[1]][r[2]][r[3]][r[4]][r[5]]+=sum;
	}
	FOR(i,1000000) {
		x=i;
		int dup=1;
		FOR(j,6) {
			r[j]=x%10;
			if(r[j]>=5) dup=0;
			x/=10;
			r[j]=9-r[j];
		}
		if(dup) ret-=1LL*C[i];
		ret+=1LL*C[i]*S[r[0]][r[1]][r[2]][r[3]][r[4]][r[5]];
		
		
	}
	cout<<ret/2<<endl;
	
}

まとめ

多次元配列の累積和は、高速ゼータ変換の要領で求めると楽なの、言われてみればそうだな。