P3047 [USACO12FEB]Nearby Cows G 题解

P3047 [USACO12FEB]Nearby Cows G

题目描述

思路

使用换根DP，

设 \(dp[i][j]\) 表示以 \(i\) 为根节点的子树中深度小于等于 \(j\) 的点的权值之和。

设 \(f[i][j]\) 表示将第 \(i\) 个点作为整棵树的根节点深度小于等于 \(j\) 的点的权值之和。

有：

\[\begin{cases}
dp[u][k] = \sum dp[to][k - 1] \\
f[to][j] = dp[to][j] - dp[to][j - 2] + f[u][j - 1]
\end{cases}
\]

\(f[to][j] = dp[to][j] - dp[to][j - 2] + f[u][j - 1]\) 表示：

代码：

#include <bits/stdc++.h>

using namespace std;

const int N = 100010, M = 30;

int w[N];

struct Edge {
	int to;
	int next;
}e[N * 2];

int head[N], idx;

void add(int a, int b) {
	idx++;
	e[idx].to = b;
	e[idx].next = head[a];
	head[a] = idx;
}

int n, k;
int dp[N][M];
int f[N][M];

void dfs(int u, int fa) {
    for (int i = 0; i <= k; i++) dp[u][i] = w[u];
    for (int i = head[u]; i; i = e[i].next) {
        int to = e[i].to;
        if (to == fa) continue;
        dfs(to, u);
        for (int i = 1; i <= k; i++) dp[u][i] += dp[to][i - 1];
    }
}

void dfs2(int u, int fa) {
    for (int i = head[u]; i; i = e[i].next) {
        int to = e[i].to;
        if (to == fa) continue;
        f[to][1] = dp[to][1] + w[u];
        for (int j = k; j >= 2; j--) f[to][j] = f[u][j - 1] - dp[to][j - 2] + dp[to][j];
        dfs2(to, u);
    }
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);

    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int a, b;
        cin >> a >> b;
        add(a, b);
        add(b, a);
    }
    for (int i = 1; i <= n; i++) cin >> w[i];
    dfs(1, 0);
    memcpy(f[1], dp[1], sizeof(f[1]));
    dfs2(1, 0);
    for (int i = 1; i <= n; i++) cout << f[i][k] << '\n';
	return 0;
}