Contest Duration: - (local time) (100 minutes) Back to Home

## E - Red and Blue Tree Editorial by en_translator

First, we count how many times the piece passes each edge with algorithm like DFS (Depth-First Search). If it passes Edge $$i$$ for $$C_i$$ times, then the problem is equivalent to “how many ways are there to separate $$C_1,\ldots,C_{N-1}$$ into two groups so that the sum of the former $$R$$ and the sum of the latter $$B$$ satisfies $$R-B=K$$?”

Moreover, with $$S=C_1+\ldots+C_{N-1}$$, it is boiled down to “how many ways are there to choose some elements from $$C_1,\ldots,C_{N-1}$$ so that its sum becomes $$\frac{K+S}{2}$$?”

Therefore, this is computed by a DP where

$$DP[i][j]=$$ The number of ways to choose some elements from $$C_1,\ldots,C_i$$ so that its sum becomes $$j$$.

Note that in some cases $$\frac{K+S}{2}$$ may not be an integer or falls below $$0$$.

The DFS for finding $$C_i$$ costs $$O(NM)$$ time, and the DP costs $$O(N(NM+|K|))$$ time, so the overall time complexity is $$O(N^2M+N|K|)$$.

Sample code (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: