AtCoder ABC #259 : Ex - Yet Another Path Counting

最近のExにしては実装軽め。
https://atcoder.jp/contests/abc259/tasks/abc259_h

問題

N*Nのグリッドがあり、各セルにラベルが振られている。
あるセルから、右または下へのセルへの移動を繰り返すパスを考える。
始点と終点のセルのラベルが同じようなパスは何通りか。

解法

始点と終点さえ定めれば、パスの組み合わせは（前処理済みの二項係数を用いれば)O(1)で求められる。
ただし、同じラベルは最大O(N^2)個あり、愚直に始点と終点の組み合わせを列挙するとO(N^4)でTLEする。

一方、典型的なグリッドを端から埋めて行くDPで、1つのラベルについてO(N^2)でまとめて処理することも可能である。

そこで、ラベルの登場頻度mに応じて両者を使い分けよう。

• m≦Nの場合、始点・終点を総当たりしてO(m^2)で数え上げよう。
• m＞Nの場合、DPでO(N^2)かけて数えよう。

m＞NとなるラベルはN個未満なので、全ラベル合計でもO(N^3)で済む。

int N;
int A[404][404];
const ll mo=998244353;
vector<pair<int,int>> cand[404*404];
ll dp[404][404];

ll comb(int P_,int Q_) {
	static const int N_=1020;
	static ll C_[N_][N_];
	
	if(C_[0][0]==0) {
		int i,j;
		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])%mo;
	}
	if(P_<0 || P_>N_ || Q_<0 || Q_>P_) return 0;
	return C_[P_][Q_];
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(y,N) FOR(x,N) {
		cin>>A[y][x];
		cand[A[y][x]].push_back({y,x});
	}
	
	ll ret=0;
	FOR(i,404*402) {
		if(cand[i].size()>=400) {
			ZERO(dp);
			FOR(y,N) FOR(x,N) {
				if(x) dp[y][x]+=dp[y][x-1];
				if(y) dp[y][x]+=dp[y-1][x];
				if(A[y][x]==i) dp[y][x]++;
				dp[y][x]%=mo;
				if(A[y][x]==i) ret+=dp[y][x];
			}
		}
		else {
			FORR2(x1,y1,cand[i]) FORR2(x2,y2,cand[i]) {
				if(x1<=x2&&y1<=y2) ret+=comb(x2+y2-x1-y1,x2-x1);
			}
			
		}
	}
	cout<<ret%mo<<endl;
}

まとめ

これは割とすんなり思いつけて良かった。