abc222_e Red and Blue Tree(树上差分+01背包)

题目链接

题目大意

  给你一棵树，和一个序列，序列描述了在树上走的路径\((a_{i-1} -> a_i)\)，你可以给一条边染成红色或者蓝色，
问走过的红色边的条数-蓝色边的条数等于K的方案数。

解题思路

  先将式子转换一下，设sum表示所有边的经过次数，那么\(R+B=sum, R-B=K\)，可以得到\(R = \frac{sum-K}{2}\)，很明显\(sum-K\)得是偶数并且非负。每条边的经过次数可以直接暴力算出来，把边转换成点，每个点的权值表示这个点的父亲与之相连的边路过的次数，当然也可以树上差分(写起来很长，也没必要，但是顺道练练)。然后每条边相当于一个物品，求01背包刚好装成R时的方案数就行了。

代码

const int maxn = 1e3+10;
const int maxm = 1e5+10;
vector<int> e[maxn];
int arr[maxn], sub[maxn], dep[maxn], f[maxn][20];
void dfs(int u) {
    for (auto v : e[u]) {
        if (dep[v]) continue;
        dep[v] = dep[u]+1;
        f[v][0] = u;
        for (int i = 1; i<20; ++i) f[v][i] = f[f[v][i-1]][i-1];
        dfs(v);
    }
}
int lca(int x, int y) {
    if (dep[x]<dep[y]) swap(x, y);
    for (int i = 19; i>=0; --i)
        if (dep[f[x][i]]>=dep[y]) x = f[x][i];
    if (x==y) return x;
    for (int i = 19; i>=0; --i)
        if (f[x][i]!=f[y][i]) x = f[x][i], y = f[y][i];
    return f[x][0];
}
vector<int> w;
int sum, fx = 1;
void dfs2(int u, int p) {
    for (auto v : e[u]) {
        if (v==p) continue;
        dfs2(v, u);
        sub[u] += sub[v];
    }
    if (sub[u]) w.push_back(sub[u]);
    else if (u!=1) fx = fx*2%MOD;
    sum += sub[u];
}
int dp[maxm];
int main() {
    IOS;
    int n, m, k; cin >> n >> m >> k;
    for (int i = 1; i<=m; ++i) cin >> arr[i];
    for (int i = 1, a, b; i<n; ++i) {
        cin >> a >> b;
        e[a].push_back(b);
        e[b].push_back(a);
    }
    dep[1] = 1;
    dfs(1);
    for (int i = 2; i<=m; ++i) {
        ++sub[arr[i-1]];
        ++sub[arr[i]];
        sub[lca(arr[i-1], arr[i])] -= 2;
    }
    dfs2(1, -1);
    if ((sum+k)%2 || sum+k<0 || sum-k<0) {
        cout << 0 << endl;
        return 0;
    }
    sum = min(sum+k, sum-k);
    sum /= 2;
    dp[0] = 1;
    for (auto v : w)
        for (int i = sum; i>=v; --i)
            dp[i] = (dp[i-v]+dp[i])%MOD;
    cout << 1LL*dp[sum]*fx%MOD << endl;
    return 0;
}