TopCoder SRM 854 : Div1 Medium SimilarNumberPairs

こちらは特にひねりもなくてストレート。
https://community.topcoder.com/stat?c=problem_statement&pm=18152

問題

6桁の整数がN個与えられる。
6桁に満たない整数は、leading zeroを含めて6桁と考えてよい。

これらの整数のペアのうち、ちょうどS桁が一致しているようなものは何通りか。

解法

包除原理で解く。
f(n) := n桁が一致している整数対の数(他の桁は一致してもしなくても良い)
g(n) := ちょうどn桁が一致している整数対の数

とすると、
g(n) = f(n) - C(6,n+1)*g(n+1) - C(6,n+2)*g(n+2) - .... - C(6,6)*f(6)
で求めることができる。

ll A[202020];

ll C[11*11*11*11*11*11];
ll p10[7],p11[7];

ll ret[7];
static const int N_=10;
static ll C_[N_][N_];

class SimilarNumberPairs {
	public:
	long long count(int N, int S, vector <int> Aprefix, int seed) {
		int i,j;
		FOR(i,Aprefix.size()) A[i]=Aprefix[i];
		ll state = seed;
		for(i=Aprefix.size();i<N;i++) {
		    state = (state * 1103515245 + 12345) %(1LL<<31);
			A[i] = (state % 1000000);
		}
		p10[0]=p11[0]=1;
		FOR(i,6) p10[i+1]=p10[i]*10,p11[i+1]=p11[i]*11;
		
	
		FOR(i,N_) C_[i][0]=C_[i][i]=1;
		for(i=1;i<N_;i++) for(j=1;j<i;j++) C_[i][j]=(C_[i-1][j-1]+C_[i-1][j]);

		ZERO(C);
		FOR(i,N) {
			int mask;
			FOR(mask,1<<6) {
				ll v=0;
				FOR(j,6) {
					if(mask&(1<<j)) v=v*11+10;
					else v=v*11+A[i]/p10[j]%10;
				}
				//cout<<mask<<" "<<A[i]<<" "<<v<<endl;
				C[v]++;
			}
		}
		
		ZERO(ret);
		FOR(i,p11[6]) {
			int match=0;
			FOR(j,6) if(i/p11[j]%11!=10) match++;
			ret[match]+=1LL*C[i]*(C[i]-1)/2;
		}
		for(i=6;i>=0;i--) {
			FOR(j,i) ret[j]-=ret[i]*C_[i][j];
		}
		
		ll sum=0;
		FOR(i,7) {
			cout<<ret[i]<<endl;
			sum+=ret[i];
		}
		assert(sum==1LL*N*(N-1)/2);
		return ret[S];
		
		
		
	}

まとめ

もっと短く書けないかな。