5898번: Nearby Cows - 최적화 필요

https://www.acmicpc.net/problem/5898

#include<stdio.h>
#include<vector>
using namespace std;
const int MAX_N = 100000;
bool ck[MAX_N + 1];
vector<int> adj[MAX_N + 1], vis;
int n, k, dp[MAX_N + 1][21], c[MAX_N + 1];
void dfs(int h) {
    ck[h] = true;
    dp[h][0] = c[h];
    vis.push_back(h);
    for (auto t : adj[h]) {
        if (ck[t]) continue;
        dfs(t);
        for (int i = 1; i <= k; i++)
            dp[h][i] += dp[t][i - 1];
    }
}
int main() {
    scanf("%d %d", &n, &k);
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        scanf("%d %d", &a, &b);
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    for (int i = 1; i <= n; i++)
        scanf("%d", &c[i]);
    dfs(1);
    ck[1] = false;
    for (int i = 1; i < n; i++) {
        int h = vis[i];
        ck[h] = false;
        for (int i = k; i >= 2; i--)
            dp[h][i] -= dp[h][i - 2];
        for (auto t : adj[h])
            if (!ck[t])
                for (int i = 1; i <= k; i++)
                    dp[h][i] += dp[t][i - 1];
    }
    for (int i = 1; i <= n; i++) {
        int s = 0;
        for (int j = 0; j <= k; j++) s += dp[i][j];
        printf("%d\n", s);
    }
    return 0;
}