题意:求树上所有满足路径长度小于等于k的两点的对数。
思路还是很好理解的,对每一棵子树,以重心为根节点统计在不同子树且相互距离小于等于k的对数。就是实现的细节处理很多,主要都是递归处理的比较多。
#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <map>
#include <vector>
using namespace std;
const int maxn = 50000 + 10;
struct edge
{
int to, dist, next;
}edges[maxn];
int edgehead[maxn], tot;
int n, k, vis[maxn], ans, root, num;
int mx[maxn], size[maxn], dis[maxn];
void addedge(int from, int to, int dist)
{
edges[tot].dist = dist;
edges[tot].to = to;
edges[tot].next = edgehead[from];
edgehead[from] = tot++;
}
void dfssize(int u, int fa)
{
size[u] = 1;
mx[u] = 0;
for(int i = edgehead[u]; i != -1; i = edges[i].next)
{
int v = edges[i].to;
if(v != fa && !vis[v])
{
dfssize(v, u);
size[u] += size[v];
mx[u] = max(mx[u], size[v]);
}
}
}
void dfsroot(int r, int u, int fa)
{
mx[u] = max(mx[u], size[r] - size[u]);
if(mx[u] < mx[root]) root = u;
for(int i = edgehead[u]; i != -1; i = edges[i].next)
{
int v = edges[i].to;
if(v != fa && !vis[v]) dfsroot(r, v, u);
}
}
void dfsdis(int u, int d, int fa)
{
dis[num++] = d;
for(int i = edgehead[u]; i != -1; i = edges[i].next)
{
int v = edges[i].to;
if(v != fa && !vis[v]) dfsdis(v, d + edges[i].dist, u);
}
}
int calc(int u, int d)
{
int ret = 0;
num = 0;
dfsdis(u, d, 0);
sort(dis, dis + num);
int i = 0, j = num - 1;
while(i < j)
{
while(dis[i] + dis[j] > k && i < j) j--;
ret += j - i;
i++;
}
return ret;
}
void dfs(int u)
{
root = u;
dfssize(u, 0);
dfsroot(u, u, 0);
ans += calc(root, 0);
vis[root] = 1;
for(int i = edgehead[root]; i != -1; i = edges[i].next)
{
int v = edges[i].to;
if(!vis[v])
{
ans -= calc(v, edges[i].dist);
dfs(v);
}
}
}
void prework()
{
memset(vis, 0, sizeof(vis));
memset(edgehead, 0xff, sizeof(edgehead)); tot = 0;
for(int i = 0; i < n - 1; ++i)
{
int a, b, c;
scanf("%d %d %d", &a, &b, &c);
addedge(a, b, c);
addedge(b, a, c);
}
}
void solve()
{
ans = 0;
dfs(1);
printf("%d\n", ans);
}
int main()
{
while(~scanf("%d %d", &n, &k))
{
if(!n && !k) break;
prework();
solve();
}
return 0;
}