AtCoder ABC #326 (Ｓｋｙ株式会社プログラミングコンテスト2023) : G - Delivery on Tree

解き切ったのはいいけど、だいぶミスが多かった。
https://atcoder.jp/contests/abc329/tasks/abc329_g

問題

根付きの二分木が与えられる。
M個のリンゴがあり、それぞれの初期位置の頂点と、移動先の頂点が与えられる。

これらのリンゴを、移動先に移動させたい。
ここで根頂点にカゴがある。カゴを辺に沿って動かしながらリンゴを積んだり下ろしたりして、条件を達成させたい。
ただし、各辺は1往復までしか利用できない。また、カゴに同時に載せられるリンゴはK個までである。

条件を満たすカゴの移動順は何通りか。

解法

下記をDFSで葉から順にみていこう。
f(v,k) := vの親頂点からvにk個のリンゴを積んだ状態で来た時、条件を満たしたままSubtree内のリンゴを適切に処理できるvのSubtree内でのカゴの動かし方の組み合わせ

まず各リンゴについて、以下を求めておく。

• 積むのはいつか(この頂点から左の子頂点、右の子頂点、親頂点のいずれに動くときか)
• 下すのはいつか(この頂点から左の子頂点、右の子頂点、親頂点のいずれからこの頂点に来たときか)

また、リンゴの初期位置と移動先の頂点の条件より、そのLCAとなる頂点において左右どちらの子頂点から先にたどらないといけないか決まる場合がある。

あとはDFSにおいて、左→右の順でたどる場合と、右→左の順でたどる場合のそれぞれのパターンの結果を足し合わせよう。
途中カゴのリンゴの個数がKを超える場合は除外する。

vector<ll> dfs2(int cur) {
	vector<ll> ret(K+1);
	int i;
	if(E[cur].empty()) {
		FOR(i,K+1) {
			if(i>=downfin[cur]&&i-downfin[cur]+up[cur]<=K) {
				ret[i]=1;
			}
		}
	}
	else {
		if(E[cur].size()==1) {
			E[cur].push_back(N);
			lef[cur]=1;
		}
		
		int s=E[cur][0];
		int t=E[cur][1];
		vector<ll> SS=dfs2(s);
		vector<ll> TT=dfs2(t);
		
		if(ri[cur]==0) {
			//左から
			FOR(i,K+1) {
				if(i<downfin[cur]) continue;
				int tar=i-downfin[cur]+downL[cur];
				if(tar<0||tar>K) continue;
				ll a=SS[tar];
				tar+=Sstart[s]-Sfin[s];
				if(tar<0||tar>K) continue;
				tar+=-upLfin[cur]+downR[cur];
				if(tar<0||tar>K) continue;
				a=a*TT[tar]%mo;
				tar+=Sstart[t]-Sfin[t];
				if(tar<0||tar>K) continue;
				tar+=-upRfin[cur]+up[cur];
				if(tar>=0&&tar<=K) ret[i]+=a;
			}
		}
		if(lef[cur]==0) {
			//右から
			FOR(i,K+1) {
				if(i<downfin[cur]) continue;
				int tar=i-downfin[cur]+downR[cur];
				if(tar<0||tar>K) continue;
				ll a=TT[tar];
				tar+=Sstart[t]-Sfin[t];
				if(tar<0||tar>K) continue;
				tar+=-upRfin[cur]+downL[cur];
				if(tar<0||tar>K) continue;
				a=a*SS[tar]%mo;
				tar+=Sstart[s]-Sfin[s];
				if(tar<0||tar>K) continue;
				tar+=-upLfin[cur]+up[cur];
				if(tar>=0&&tar<=K) ret[i]+=a;
			}
		}
	}
	FORR(r,ret) r%=mo;
	return ret;
}

void dfs(int cur) {
	L[cur]=id++;
	Sstart[cur]=start[cur];
	Sfin[cur]=fin[cur];
	FORR(e,E[cur]) {
		dfs(e);
		Sstart[cur]+=Sstart[e];
		Sfin[cur]+=Sfin[e];
	};
	R[cur]=id++;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	FOR(i,N-1) {
		cin>>P[0][i+1];
		P[0][i+1]--;
		D[i+1]=D[P[0][i+1]]+1;
		E[P[0][i+1]].push_back(i+1);
	}
	FOR(i,20) {
		FOR(j,N) P[i+1][j]=P[i][P[i][j]];
	}
	
	FOR(i,M) {
		cin>>S[i]>>T[i];
		S[i]--,T[i]--;
		start[S[i]]++;
		fin[T[i]]++;
	}
	dfs(0);
	FOR(i,M) {
		if(L[S[i]]<L[T[i]]&&L[T[i]]<R[S[i]]) {
			if(R[E[S[i]][0]]<=L[T[i]]) {
				downR[S[i]]++;
			}
			else {
				downL[S[i]]++;
			}
			downfin[T[i]]++;
		}
		else if(L[T[i]]<L[S[i]]&&L[S[i]]<R[T[i]]) {
			up[S[i]]++;
			if(R[E[T[i]][0]]<=L[S[i]]) {
				upRfin[T[i]]++;
			}
			else {
				upLfin[T[i]]++;
			}
		}
		else {
			x=S[i],y=T[i];
			up[x]++;
			FOR(j,20) {
				if(D[x]>D[y]&&((D[x]-D[y])&(1<<j))) x=P[j][x];
				if(D[y]>D[x]&&((D[y]-D[x])&(1<<j))) y=P[j][y];
			}
			for(j=19;j>=0;j--) if(P[j][x]!=P[j][y]) {
				x=P[j][x];
				y=P[j][y];
			}
			k=P[0][x];
			if(E[k][0]==x) {
				lef[k]++;
			}
			else {
				ri[k]++;
			}
			downfin[T[i]]++;
		}
	}
	auto R=dfs2(0);
	cout<<R[0]<<endl;
	
}

まとめ

色々こんがらがって何か所もバグを作りこんでしまった、反省。