POJ1741 Tree（点分治）

传送门

题意：
给出一颗边权为正数的树，问有多少点对之间的距离小于等于(k)。

思路：
点分治模板题，对于一个子问题，(dfs)出所有的距离之后，排个序可以用双指针，或者直接二分也行，复杂度都为(O(nlogn^2))。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 40005;
struct Edge{
    int v, next, w;
}e[N << 1];
int head[N], tot;
void adde(int u, int v, int w) {
    e[tot].v = v; e[tot].w = w; e[tot].next = head[u]; head[u] = tot++;
}
int n, k;
int ans, mx, rt, num;
int sz[N], dis[N], vis[N];
void dfs_rt(int u, int fa) {
    sz[u] = 1;
    int Mx = 0;
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v;
        if(v == fa || vis[v]) continue ;
        dfs_rt(v, u) ;
        Mx = max(Mx, sz[v]) ;
        sz[u] += sz[v] ;
    }
    Mx = max(Mx, num - sz[u]) ;
    if(Mx < mx) mx = Mx, rt = u;
}
int cnt;
void dfs_dis(int u, int fa, int D) {
    dis[++cnt] = D;
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v, w = e[i].w;
        if(vis[v] || v == fa) continue ;
        dfs_dis(v, u, D + w) ;
    }
}
int calc(int u, int D) {
    cnt = 0, dfs_dis(u, 0, D) ;
    sort(dis + 1, dis + cnt + 1) ;
    int r = cnt, sum = 0;
    for(int l = 1; l <= cnt && l <= r; l++) {
        while(l < r && dis[r] + dis[l] > k) r--;
        sum += r - l;
    }
    return sum ;
}
void dfs_sz(int u, int fa) {
    sz[u] = 1;
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v;
        if(vis[v] || fa == v) continue;
        dfs_sz(v, u);
        sz[u] += sz[v] ;
    }
}
void dfs(int u) {
    ans += calc(u, 0) ;vis[u] = 1;
    dfs_sz(u, 0) ;
    for(int i = head[u]; i != -1; i = e[i].next) {
        int v = e[i].v;
        if(vis[v]) continue ;
        ans -= calc(v, e[i].w) ;
        mx = num = sz[v], dfs_rt(v, u);
        dfs(rt);
    }
}
int main() {
    memset(head, -1, sizeof(head)) ;
    cin >> n;
    for(int i = 1; i < n; i++) {
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w) ;
        adde(u, v, w); adde(v, u, w) ;
    }
    cin >> k;
    mx = num = n, dfs_rt(1, 0) ;
    dfs(rt) ;
    cout << ans;
    return 0;
}