题目描述

给定一个 个点的树，树有点权。

如果 为 和 的最近公共祖先(LCA)，并且 和 的树上距离等于 ，那么点 的答案就会加上 。

求所有点的答案。

正解

考虑暴力 dp，设 表示 的子树内离 的距离为 的点的个数，设 表示 的子树内离 的距离为 的点权和。

每次合并一颗子树，更新答案并且更新 dp 数组，时间复杂度 。

发现 dp 复杂度只跟深度有关，那么长链剖分后可以优化复杂度至 。

其实还有 dsu 的做法，每次暴力合并复杂度也没有问题，是 的。

这里给出长链剖分的代码 :

#include <bits/stdc++.h>
#define N 100005

using namespace std;

int n, k;
int o[N], d[N], a[N], totVec;
int fa[N], wson[N], dist[N];
int head[N], nex[N << 1], to[N << 1], ecnt;

long long ans[N];
vector<long long> f[N], g[N];


inline int read() {
    int x = 0; char ch = getchar();
    while(!isdigit(ch)) ch = getchar();
    while(isdigit(ch)) x = x * 10 + ch - '0', ch = getchar();
    return x;
}
inline void addE(int u, int v) {
    to[++ecnt] = v;
    nex[ecnt] = head[u], head[u] = ecnt;
}

void preDfs(int u) {
    for(int i = head[u], v; i; i = nex[i]) {
        v = to[i];
        if(v == fa[u]) continue;
        fa[v] = u, preDfs(v);
        if(dist[wson[u]] < dist[v])
            wson[u] = v;
    }
    dist[u] = dist[wson[u]] + 1;
}

void dfs(int u, int top) {
    if(u == top) {
        o[u] = ++totVec;
        f[o[u]].resize(dist[u]);
        g[o[u]].resize(dist[u]);
    }
    if(wson[u]) {
        o[wson[u]] = o[u];
        d[wson[u]] = d[u] + 1;
        dfs(wson[u], top);
    }

    ++f[o[u]][d[u]];
    g[o[u]][d[u]] += a[u];
    for(int i = head[u], v; i; i = nex[i]) {
        v = to[i];
        if(v == fa[u] || v == wson[u]) continue;
        dfs(v, v);
        for(int j = 0; j < dist[v]; ++j)
            if(k - j - 1 > 0 && k - j - 1 + d[u] < dist[top]) {
                ans[u] += f[o[u]][k - j - 1 + d[u]] * g[o[v]][j];
                ans[u] += g[o[u]][k - j - 1 + d[u]] * f[o[v]][j];
            }
        for(int j = 0; j < dist[v]; ++j) {
            f[o[u]][j + 1 + d[u]] += f[o[v]][j];
            g[o[u]][j + 1 + d[u]] += g[o[v]][j];
        }
    }
}

int main() {
    n = read(), k = read();
    for(int i = 1; i <= n; ++i)
        a[i] = read();
    for(int i = 1, u, v; i < n; ++i) {
        u = read(), v = read();
        addE(u, v), addE(v, u);
    }

    preDfs(1);
    dfs(1, 1);

    for(int i = 1; i <= n; ++i)
        printf("%lld%c", ans[i], " \n"[i == n]);
    return 0;
}