Tree

luogu 4178

金牌导航 点分治-1

题目大意

给出一棵树，问你书中路径长度小于等于k的点对个数有多少个

输入样例

5
1 2 3
1 3 1
1 4 2
3 5 1
4

输出样例

8

数据范围

$1⩽N⩽4\times 10^4$

解题思路

对于该树，求出dfs求出重心
对于两点都在同一子树中的点对，分治处理即可
而对于两点不在同一子树的
可以先求出重心到所有点的距离，然后排个序，用双指针得出长度小于k的数量
但是这样得出来的点对可能在同一子树中（图如下）
对于这种情况，我们计算一遍子树的，然后减去就行了（对于在2的同一子树中的情况，2中有出现，1中也有出现，所以2就不用再减子树的点对数了）

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define ll long long
#define N 40010
using namespace std;
int n, k, x, y, z, l, r, h, qn, sum, tot, ans, s[N], p[N], q[N], mx[N], size[N], head[N];
struct rec
{int to, l, next;
}a[N<<1];
void add(int x, int y, int z)
{a[++tot].to = y;a[tot].l = z;;a[tot].next = head[x];head[x] = tot;
}
void dfs1(int x)//求重心
{p[x] = 1;size[x] = 1;for (int i = head[x]; i; i = a[i].next)if (!p[a[i].to]){dfs1(a[i].to);size[x] += size[a[i].to];mx[x] = max(mx[x], size[a[i].to]);}p[x] = 0;mx[x] = max(mx[x], sum - size[x]);//还要计算除去x子树后其它点的数量if (mx[x] <= mx[h]) h = x;//求重心return;
}
void dfs2(int x)
{q[++qn] = s[x];//求出到根节点的最短距离p[x] = 1;for (int i = head[x]; i; i = a[i].next)if (!p[a[i].to]){s[a[i].to] = s[x] + a[i].l;dfs2(a[i].to);}p[x] = 0;
}
int count(int x, int y)
{int sum = 0;s[x] = y;qn = 0;dfs2(x);sort(q + 1, q + 1 + qn);l = 1;r = qn;while(l < r){while(q[l] + q[r] > k && l < r) r--;//如果大于k，那么r--sum += r - l;//计算答案l++;//如果l变大了，那r只可能变小}return sum;
}
void solve(int x)
{h = 0;mx[0] = n;dfs1(x);int G = h;ans += count(G, 0);p[G] = 1;for (int i = head[G]; i; i = a[i].next)if (!p[a[i].to]) ans -= count(a[i].to, a[i].l);//减去矛盾的for (int i = head[G]; i; i = a[i].next)if (!p[a[i].to]) sum = size[a[i].to], solve(a[i].to);//分治
}
int main()
{scanf("%d", &n);for (int i = 1; i < n; ++i){scanf("%d%d%d", &x, &y, &z);add(x, y, z);add(y, x, z);		}scanf("%d", &k);sum = n;solve(1);printf("%d", ans);return 0;
}