Official
E - Red and Blue Tree Editorial by kyopro_friends
まず、駒の移動により各辺を何度通過するかをDFSなどを用いて求めます。辺 \(i\) を \(C_i\) 回通過するとき、問題は「\(C_1,\ldots,C_{N-1}\) を\(2\) 組に分け、一方の合計を \(R\)、他方の合計を \(B\) とするとき、\(R-B=K\) とする方法は何通りか?」となります。
さらに、\(S=C_1+\ldots+C_{N-1}\) とおくと、「\(C_1,\ldots,C_{N-1}\) からいくつか選んで、和を \(\frac{K+S}{2}\) にする方法は何通りか?」と読み替えることができます。
したがってこれは
\(DP[i][j]=\) \(C_1,\ldots,C_i\) からいくつか選んで和を \(j\) にする方法の数
というDPにより求めることができます。
\(\frac{K+S}{2}\) が整数にならないケースや \(0\) 未満になるケースに注意してください。
\(C_i\) を求めるために DFS をする部分が \(O(NM)\) 、DP部分が \(O(N(NM+|K|))\) かかり、計算量は全体で \(O(N^2M+N|K|)\) となります。
実装例(C++)
#include<bits/stdc++.h>
using namespace std;
const int mod=998244353;
vector<vector<pair<int,int>>>G;
vector<int>C;
int dfs(int v,int pre,int goal){
if(v==goal)return 1;
for(auto [vv,i]:G[v])if(vv!=pre){
if(dfs(vv,v,goal)){
C[i]++;
return 1;
}
}
return 0;
}
int main(){
int N, M, K;
cin >> N >> M >> K;
vector<int>A(M);
for(int i=0;i<M;i++){
cin >> A[i];
A[i]--;
}
G.resize(N);
for(int i=0;i<N-1;i++){
int x,y;
cin >> x >> y;
x--,y--;
G[x].push_back({y,i});
G[y].push_back({x,i});
}
C.resize(N-1);
for(int i=0;i<M-1;i++)dfs(A[i],-1,A[i+1]);
int S=0;
for(int i=0;i<N-1;i++)S+=C[i];
if((S+K)%2 || S+K<0){
cout << 0;
return 0;
}
vector<int>dp(100001);
dp[0]=1;
for(int i=0;i<N-1;i++)for(int x=100000;x>=C[i];x--)(dp[x]+=dp[x-C[i]])%=mod;
cout << dp[(S+K)/2];
}
posted:
last update: