kmjp's blog

競技プログラミング参加記です

AtCoder ABC #160 : F - Distributing Integers

実装部分はともかく、考察で迷わない分こっちの方が相対的には得意だなぁ。
https://atcoder.jp/contests/abc160/tasks/abc160_f

問題

木を成すグラフが与えられる。
ある点を根としたとき、その頂点に数値1を振る。
その後、数値設定済みの隣接点のいずれか未設定の点に、未登場の最少の正整数を振る、ということを順次行う。

各頂点を根としたときの、整数の振り方は何通りか。

解法

典型的な木DP。
ある木のSubTreeについて考えると、今その頂点vのSubTreeにC(v)個の点があり、そのC(v)点に1~C(v)の値を振るやり方がD(v)通りあるとする。
ここに新たなSubTree wをvの子に連結した状態をv'とする。
そうするとC(v')=C(v)+C(w)でD(v')=D(v)*D(w)*Comb(C(v)+C(w)-1,C(w))となる。
これはvのSubtreeのうちv以外の(C(v)-1)頂点に振った値とw以下のC(w)頂点に振った値を改めて2~(C(v)+C(w))番に振り分けるやり方を考えることに相当する。

また上記の演算では逆に1個子頂点(とそのSubTree)をはがすときの計算もできる。
なので全方位木DPの要領で、親方向も含めてD(v)を求めることができる。

int N;
vector<int> E[202020];
ll dp[202020];
int C[202020];
ll pdp[202020];
const ll mo=1000000007;

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

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

void dfs(int cur,int pre) {
	C[cur]=0;
	dp[cur]=1;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		
		dp[cur]=dp[cur]*dp[e]%mo*comb(C[cur]+C[e],C[e])%mo;
		C[cur]+=C[e];
	}
	C[cur]++;
}

void dfs2(int cur,int pre,ll pv) {
	dp[cur]=dp[cur]*pv%mo*comb(N-1,N-C[cur])%mo;
	FORR(e,E[cur]) if(e!=pre) {
		
		pv=dp[cur]*modpow(dp[e]*comb(N-1,C[e])%mo)%mo;
		dfs2(e,cur,pv);
		
		
		
	}
	
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,0);
	dfs2(0,0,1);
	
	FOR(i,N) cout<<dp[i]<<endl;
}

まとめ

この問題、典型っぽいしどこかで出てそうな気はする。