AtCoder ARC #140 : D - One to One

Cまでは良かったんだが。
https://atcoder.jp/contests/arc140/tasks/arc140_d

問題

各要素1～Nの値を取る、N要素の整数列Xに対し、以下の値をf(X)とする。

• N頂点の無向グラフにおいて、頂点iとX[i]をそれぞれ結んだ時、連結成分の個数

今ここで、同じく各要素1～Nの値を取る、N要素の整数列Aが与えられる。
ただし、一部要素は値が未定である。

Aのうち未定の箇所がそれぞれ1～Nの値を取った場合、得られる数列Xにおけるf(X)の総和を求めよ。

解法

Aのうち確定状態の要素について、辺を張ると、辺の数=頂点の数である連結成分と、辺の数+1=頂点の数である連結成分ができる。
前者は、未定の値に対応する頂点を含まず、後者は未定の値に対応する頂点を１つ含んでいる。
前者はもう連結成分１個できるのが確定なので、(連結成分数)*(N^(未定の成分数))を解に計上しよう。

あとは、未定の値に対応する頂点が、どのように辺の数=頂点の数となる連結成分を新規に構築するかを考える。
もし各(辺の数+1=頂点の数である)連結成分がいくつか集まって(辺の数=頂点の数である)連結成分を作る場合、それぞれV_i個の頂点を持つk個の連結成分が、さらに連結する場合、その組み合わせはProd(V)*(k-1)!である。

よって、
dp(n,k) := 連結成分中のうちprefix n個の中から、(辺の数=頂点の数である)連結成分を作るためにk個を選んだ時のprod(V)の総和
をDPで求めよう。
あとはdp(未定の成分数,k) * (N^(未定の成分数-k))の総和を解に計上すればよい。

int N,A[2020];
const ll mo=998244353;

template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit(int num=um) {int i; FOR(i,num) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<2020> uf;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll dp[2020];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	int unf=0;
	FOR(i,N) {
		cin>>A[i];
		if(A[i]>0) {
			A[i]--;
			uf(i,A[i]);
		}
		else {
			unf++;
		}
	}
	ll ret=0;

	dp[0]=1;
	FOR(i,N) if(uf[i]==i) {
		int fix=1;
		FOR(j,N) if(uf[j]==i) {
			if(A[j]==-1) fix=0;
		}
		
		if(fix) {
			ret+=modpow(N,unf);
		}
		else {
			for(x=unf;x>=0;x--) if(dp[x]) {
				(dp[x+1]+=dp[x]*uf.count(i))%=mo;
			}
		}
	}
	
	ll f=1;
	for(i=1;i<=unf;i++) if(dp[i]) {
		if(i>1) f=f*(i-1)%mo;
		ret+=f*dp[i]%mo*modpow(N,unf-i)%mo;
	}
	
	
	
	cout<<ret%mo<<endl;
}

まとめ

本番もっと遠回りなコードを書いてしまった。