kmjp's blog

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

AtCoder ARC #121 (NOMURA プログラミングコンテスト 2021) : E - Directed Tree

土壇場で解けてギリギリ大怪我を防いだ。
https://atcoder.jp/contests/arc121/tasks/arc121_e

問題

1~N番の番号が振られたN頂点からなる根付き木が与えられる。
1番の頂点が根であり、根に近い方が頂点番号が小さい。

各頂点に1~Nの値のいずれかを、同じ値を1回ずつ割り当てることを考える。
A[i]をiが割り当てられた頂点とする。
この時、A[i]が割り当てられた頂点から、根に向けて1回以上辺に沿って移動しても、i番の頂点を通らないようなAはN!通りのうち何通りか。

解法

以下の値を考える。
f(v,n) := 頂点vのSubTree内において、条件に違反することが確定する要素数がn個であるような組み合わせの数
これを子から順に求めて行こう。
頂点vに、vのSubTree内の要素を割り当ててしまうと、nが1増える。

あとはf(1,*)を包除原理の要領で加減算していく。

int N;
int P[2020];
vector<int> E[2020];
int C[2020],D[2020];
const ll mo=998244353;
ll dp[2020][2020];

const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll pat[2020][2020];

ll comb(ll N_, ll C_) {
	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 hcomb(int P_,int Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}

int dfs(int cur,int d) {
	C[cur]=1;
	
	ll from[2020]={};
	from[0]=1;
	// ok
	FORR(e,E[cur]) {
		dfs(e,d+1);
		
		int x,y,z,a,b;
		ll to[2020]={};
		FOR(x,C[cur]+1) FOR(y,C[e]+1) {
			(to[x+y]+=from[x]*dp[e][y])%=mo;
		}
		
		swap(from,to);
		C[cur]+=C[e];
	}
	
	int x;
	FOR(x,C[cur]) {
		(dp[cur][x]+=from[x])%=mo; //safe
		(dp[cur][x+1]+=(C[cur]-(x+1))*from[x])%=mo; // ng
	}
	
	
	return C[cur];
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	comb(0,0);
	
	cin>>N;
	FOR(i,N-1) {
		cin>>P[i+1];
		P[i+1]--;
		E[P[i+1]].push_back(i+1);
	}
	
	dfs(0,1);
	
	ll ret=0;
	FOR(i,N) {
		if(i%2==0) (ret+=dp[0][i]*fact[N-i])%=mo;
		else (ret-=dp[0][i]*fact[N-i])%=mo;
	}
	
	cout<<(ret+mo)%mo<<endl;
	
}

まとめ

もっとサラッと解きたいな。