POJ1741 点分治模板

传送门：http://poj.org/problem?id=1741

题意：

求树上两点间路径长度小于k的点对个数

题解：

参考资料

守望的淀粉质略解：https://www.luogu.org/blog/user9012/dian-fen-zhi-lve-xie

粉红兔大佬的淀粉质:https://www.cnblogs.com/PinkRabbit/p/8593080.html

算法步骤

1. 计算重心位置
2. 计算答案
3. 分治子问题继续求解（1~3）

代码：

#include <set>
#include <map>
#include <cmath>
#include <cstdio>
#include <string>
#include <vector>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
typedef long long LL;
typedef pair<int, int> pii;
typedef unsigned long long uLL;
#define ls rt<<1
#define rs rt<<1|1
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
#define bug printf("*********\n")
#define FIN freopen("input.txt","r",stdin);
#define FON freopen("output.txt","w+",stdout);
#define IO ios::sync_with_stdio(false),cin.tie(0)
#define debug1(x) cout<<"["<<#x<<" "<<(x)<<"]\n"
#define debug2(x,y) cout<<"["<<#x<<" "<<(x)<<" "<<#y<<" "<<(y)<<"]\n"
#define debug3(x,y,z) cout<<"["<<#x<<" "<<(x)<<" "<<#y<<" "<<(y)<<" "<<#z<<" "<<z<<"]\n"
const int maxn = 3e5 + 5;
const int INF = 0x3f3f3f3f;
int n, k, Ans;

struct EDGE {
    int v, nxt, w;
} edge[maxn << 1];
int head[maxn], tot;
void add_edge(int u, int v, int w) {
    edge[tot].v = v;
    edge[tot].w = w;
    edge[tot].nxt = head[u];
    head[u] = tot++;
}
bool vis[10001];
int Root, Tsiz, siz[10001], wt[10001];
int arr[10001], cnt;

void GetRoot(int u, int f) {
    siz[u] = 1; wt[u] = 0;
    for(int i = head[u]; i != -1; i = edge[i].nxt) {
        int v = edge[i].v;
        if(v != f && !vis[v]) {
            GetRoot(v, u);
            siz[u] += siz[v];
            wt[u] = max(wt[u], siz[v]);
        }
    }
    wt[u] = max(wt[u], Tsiz - siz[u]);
    if(wt[Root] > wt[u]) Root = u;
}

void Dfs(int u, int D, int f) {
    arr[++cnt] = D;
    for(int i = head[u]; i != -1; i = edge[i].nxt) {
        int v = edge[i].v;
        if(v != f && !vis[v]) {
            Dfs(v, D + edge[i].w, u);
        }
    }
}

int calc(int u, int D) {
    cnt = 0; Dfs(u, D, 0); int l = 1, r = cnt, sum = 0;
    sort(arr + 1, arr + cnt + 1);
    for(;; ++l) {
        while(r && arr[l] + arr[r] > k) --r;
        if(r < l) break;
        sum += r - l + 1;
    }
    return sum;
}

void DFS(int u) {
    Ans += calc(u, 0); vis[u] = 1;
    for(int i = head[u]; i != -1; i = edge[i].nxt) {
        int v = edge[i].v;
        if(!vis[v]) {
            Ans -= calc(v, edge[i].w);
            Root = 0, Tsiz = siz[v], GetRoot(v, 0);
            DFS(Root);
        }
    }
}

int main() {
#ifndef ONLINE_JUDGE
    FIN
#endif
    while(~scanf("%d%d", &n, &k) && n && k) {
        tot = Ans = 0;
        memset(vis, 0, sizeof(vis));
        memset(head, -1, sizeof(head));
        for(int i = 2, u, v, w; i <= n; i++) {
            scanf("%d%d%d", &u, &v, &w);
            add_edge(u, v, w);
            add_edge(v, u, w);
        }
        wt[0] = INF;
        Tsiz = n;
        GetRoot(1, 0);
        DFS(Root);
        printf("%d\n", Ans - n);
    }
    return 0;
}