Solution
我们可以把边与边的到达点建一个联系,用点来表示会简单得多。
发现每条边的重要度可以 \(\mathtt{O(n)}\) 计算。那么我们就是计算每个点在 k 距离里的点的权值和。(严格来说这并不准确,然而我没找到更好的叙述方法)
我们令 \(\mathtt{dp[i][j]}\) 为 i 在 j 距离内的子树中的点的权值和。(注意这里每个点的权值是自己儿子与自己相连的边的权值和,因为一个点施工,与它相连的边都会遭殃)这个用树形 DP 应该是 \(\mathtt{O(n*k)}\) 的。
最后我们再计算每一个点的长辈与兄弟。因为父亲的权值会包含自己的,所以减去自己那一部分就好了。
注意最后一个点要计算与自己相关的点(还要包含与父亲相连的那条边)。
Code
#include <cstdio>
#include <iostream>
using namespace std;
typedef long long ll;
const int N = 3e4 + 5;
int n, k, f[N], head[N], dot[N << 1], nxt[N << 1], cnt;
ll ans, dp[N][205], siz[N], val[N], s[N];
int read() {
int x = 0, f = 1; char s;
while((s = getchar()) < '0' || s > '9') if(s == '-') f = -1;
while(s >= '0' && s <= '9') {x = (x << 1) + (x << 3) + (s ^ 48); s = getchar();}
return x * f;
}
void addEdge(const int u, const int v) {
dot[++ cnt] = v; nxt[cnt] = head[u]; head[u] = cnt;
}
void init(const int u, const int fa) {
siz[u] = 1; f[u] = fa;
for(int i = head[u]; i; i = nxt[i]) {
int v = dot[i];
if(v == fa) continue;
init(v, u);
val[v] = siz[v] * (n - siz[v]);
s[u] += val[v];
siz[u] += siz[v];
}
}
void cal(const int u) {
for(int i = 0; i <= k; ++ i) dp[u][i] = s[u];
for(int i = head[u]; i; i = nxt[i]) {
int v = dot[i];
if(v == f[u]) continue;
cal(v);
for(int j = 1; j <= k; ++ j) dp[u][j] += dp[v][j - 1];
}
}
ll get(int u) {
ll res = dp[u][k]; int fa;
for(int i = 1; i <= k; ++ i) {
fa = f[u];
if(! fa) continue;
if(i == k) {res += s[fa]; u = fa; break;}
res += dp[fa][k - i] - dp[u][k - i - 1];
u = fa;
}
return res + val[u];
}
int main() {
int u, v;
n = read(), k = read();
for(int i = 1; i < n; ++ i) {
u = read(), v = read();
addEdge(u, v), addEdge(v, u);
}
init(1, 0); cal(1);
for(int i = 1; i <= n; ++ i) ans = max(ans, get(i));
printf("%lld\n", ans);
return 0;
}