ABC222 E - Red and Blue Tree（思维+dp）

题意：

解法：

O(nm)预处理每条边被经过的次数.

题目要求R-B=K
那么问题变为n-1个数cnt[]染色,满足上面的条件.

令d[i][j]表示前i条边,R的总和为j的方案数,这个可以dp计算.
由于S=sum(cnt[])是固定的,因此知道R就能计算出B.

最后枚举R=i,如果 i-(S-i)==k,那么ans+=d[n][i].

为了防止爆空间,dp数组的第一维我优化掉了.

code：

#include<bits/stdc++.h>
// #define MULTI_CASE
#define SYNC_OFF
#define PI pair<int,int>
#define ll long long
// #define int long long
using namespace std;
const int mod=998244353;
// const int mod=1e9+7;
const int maxm=1e3+5;
vector<PI>g[maxm];
int cnt[maxm];
int d[111111];
int a[maxm];
int n,m,k;
int dfs(int x,int fa,int y){
    if(x==y)return 1;
    for(auto i:g[x]){
        int v=i.first;
        if(v==fa)continue;
        if(dfs(v,x,y)){
            cnt[i.second]++;
            return 1;
        }
    }
    return 0;
}
void solve(){
    cin>>n>>m>>k;
    for(int i=1;i<=m;i++){
        cin>>a[i];
    }
    for(int i=1;i<n;i++){
        int x,y;cin>>x>>y;
        g[x].push_back({y,i});
        g[y].push_back({x,i});
    }
    //预处理每条边被经过多少次
    for(int i=1;i<m;i++){
        dfs(a[i],-1,a[i+1]);
    }
    //dp求方案数
    //d[i]表示R=i的方案数.
    int sum=0;
    for(int i=1;i<n;i++){
        sum+=cnt[i];
    }
    d[0]=1;
    for(int i=1;i<n;i++){
        for(int j=sum;j>=cnt[i];j--){
            d[j]=(d[j]+d[j-cnt[i]])%mod;
        }
    }
    //
    int ans=0;
    for(int i=0;i<=sum;i++){
        if(i-(sum-i)==k){
            ans=(ans+d[i])%mod;
        }
    }
    ans=(ans%mod+mod)%mod;
    cout<<ans<<endl;
}
void Main(){
    #ifdef MULTI_CASE
    int T;cin>>T;while(T--)
    #endif
    solve();
}
void Init(){
    #ifdef SYNC_OFF
    ios::sync_with_stdio(0);cin.tie(0);
    #endif
    #ifndef ONLINE_JUDGE
    freopen("../in.txt","r",stdin);
    freopen("../out.txt","w",stdout);
    #endif
}
signed main(){
    Init();
    Main();
    return 0;
}

---