AtCoder ARC #107 : E - Mex Mat

これ本番に思いつく気がしないなぁ。
https://atcoder.jp/contests/arc107/tasks/arc107_e

問題

N*Nの行列を考える。
1行目と1列目は入力で与えられる。
それ以外の要素は、左隣りと上の要素のmex値を取る。

この行列内の0,1,2の登場頻度を求めよ。

解法

5列目・5行目以降の値は左上と同じ値を取る。
なので、4列目・4行目まで愚直に値を求めればよい。

int N;
int A[4][505050],B[505050][4];
ll C[3]={};

int mex[3][3]= {{1,2,1},{2,0,0},{1,0,0}};

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[0][i];
		if(i<4) B[0][i]=A[0][i];
		C[A[0][i]]++;
	}
	FOR(i,N-1) {
		cin>>B[i+1][0];
		if(i+1<4) A[i+1][0]=B[i+1][0];
		C[B[i+1][0]]++;
	}
	
	for(y=1;y<min(4,N);y++) {
		for(x=1;x<N;x++) {
			A[y][x]=mex[A[y-1][x]][A[y][x-1]];
			C[A[y][x]]++;
		}
	}
	for(x=1;x<min(4,N);x++) {
		for(y=1;y<N;y++) {
			B[y][x]=mex[B[y-1][x]][B[y][x-1]];
			if(y>=4) C[B[y][x]]++;
		}
	}
	
	if(N>=5) {
		for(i=3;i<N;i++) C[A[3][i]]+=N-1-i;
		for(i=4;i<N;i++) C[B[i][3]]+=N-1-i;
	}
	cout<<C[0]<<" "<<C[1]<<" "<<C[2]<<endl;
	
}

まとめ

これでokってどう確信持てばいいんだろうな。
実験するしかないのかな。