【换根DP+容斥】P3047 [USACO12FEB]Nearby Cows G

P3047 [USACO12FEB]Nearby Cows G - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

题意：

思路：

做法就是换根

预处理dp[v][j]用普通的树形DP处理即可

注意：一开始预处理的dp[v][j]指的是在v的子树里离v为j的权值和

Code：

#include <bits/stdc++.h>

#define int long long
#define max(a,b) (a>b?a:b)
#define min(a,b) (a<b?a:b)

using namespace std;

const int mxn=3e6+10;
const int mxe=3e5+10;
const int mod=1e9+7;

struct ty{
    int to,next;
}edge[mxe<<2];

int N,K,u,v,x,tot=0;
int head[mxn],dp[mxn][22];

void add(int u,int v){
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void G_init(){
    tot=0;
    for(int i=0;i<=N;i++){
        head[i]=-1;
    }
}
void dfs1(int u,int fa){
    for(int i=head[u];~i;i=edge[i].next){
        if(edge[i].to==fa) continue;
        dfs1(edge[i].to,u);
        for(int j=1;j<=K;j++){
            dp[u][j]+=dp[edge[i].to][j-1];
        }
    }
}
void dfs2(int u,int fa){
    for(int i=head[u];~i;i=edge[i].next){
        if(edge[i].to==fa) continue;
        for(int j=K;j>=2;j--){
            dp[edge[i].to][j]-=dp[edge[i].to][j-2];
        }
        for(int j=1;j<=K;j++){
            dp[edge[i].to][j]+=dp[u][j-1];
        }
        dfs2(edge[i].to,u);
    }
}
void solve(){
    cin>>N>>K;
    G_init();
    for(int i=1;i<=N-1;i++){
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=N;i++){
        cin>>x;
        dp[i][0]=x;
    }
    dfs1(1,0);
    dfs2(1,0);
    for(int i=1;i<=N;i++){
        int ans=0;
        for(int j=0;j<=K;j++) ans+=dp[i][j];
        cout<<ans<<'\n';
    }
}
signed main(){
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    int __=1;//cin>>__;
    while(__--)solve();return 0;
}

2023.10.7 upd：

换根原来有通用写法，学会了，严格鸽是我爹

#include <bits/stdc++.h>

#define int long long

using i64 = long long;

#define lowbit(x) (x & (- x))

constexpr int N = 1e5 + 10;
constexpr int mod = 1e9 + 7;
constexpr int Inf = 0x3f3f3f3f;
constexpr double eps = 1e-10;

int n, k;
int c[N];
int dp[N][33];
std::vector<int> adj[N];

void dfs1(int u, int fa) {
    for (auto v : adj[u]) {
        if (v == fa) continue;
        dfs1(v, u);
        for (int j = 1; j <= k; j ++) {
            dp[u][j] += dp[v][j - 1];
        }
    }
}
void dfs2(int u, int fa) {
    for (auto v : adj[u]) {
        if (v == fa) continue;
        std::vector<int> a(33);
        for (int j = 1; j <= k; j ++) {
            dp[u][j] -= dp[v][j - 1];
            a[j - 1] = dp[v][j - 1];
        }
        for (int j = 1; j <= k; j ++) {
            dp[v][j] += dp[u][j - 1];
        }
        dfs2(v, u);
        for (int j = 1; j <= k; j ++) {
            dp[u][j] += a[j - 1];
        }
    }
}
void solve() {
    std::cin >> n >> k;
    for (int i = 1; i <= n - 1; i ++) {
        int u, v;
        std::cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int i = 1; i <= n; i ++) std::cin >> c[i];
    for (int i = 1; i <= n; i ++) {
        dp[i][0] = c[i];
    }
    dfs1(1, 0);
    dfs2(1, 0);
    for (int i = 1; i <= n; i ++) {
        int ans = 0;
        for (int j = 0; j <= k; j ++) ans += dp[i][j];
        std::cout << ans << "\n";
    }
}
signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int t = 1;
    while(t --) {
        solve();
    }
    return 0;
}