点分治模板题

收藏好久的模板, 这里发一下

#include<cstdio>
#include<cctype>
#include<cstring>
#include<algorithm>
using namespace std;
inline int read()
{
    int x = 0, flag = 1;
    char c;
    while(! isgraph(c = getchar()))
        if(c == '-')
            flag *= - 1;
    while(isgraph(c))
        x = x * 10 + c - '0', c = getchar();
    x *= flag;
    return x;
}
inline void print(int x)
{
    if(x < 0)
        putchar('-');
    if(x == 0)
        putchar('0');
    int ans[10], top = 0;
    while(x)
        ans[top ++] = x % 10, x /= 10;
    for(; top; top --)
        putchar(ans[top - 1] + '0');
}
const int MAXN = (int)1e4 + (1 << 4);
int n, k;
int top;
struct edge
{
    int v, w, next;
}G[MAXN << 1];
int head[MAXN];
int vis[MAXN], mx[MAXN], dis[MAXN], size[MAXN];
void add_edge(int u, int v, int w)
{
    G[top].v = v;
    G[top].w = w;
    G[top].next = head[u];
    head[u] = top ++;
}
int ans, mn, root, num;
void get_size(int u, int fa)
{
    size[u] = 1;
    mx[u] = 0;
    for(int i = head[u]; i != - 1; i = G[i].next)
    {
        int v = G[i].v;
        if(v != fa && ! vis[v])
        {
            get_size(v, u);
            size[u] += size[v];
            if(size[v] > mx[u])
                mx[u] = size[v];
        }
    }
}
void get_root(int anc, int u, int fa)
{
    mx[u] = max(mx[u], size[anc] - size[u]);
    if(mx[u] < mn)
        mn = mx[u], root = u;
    for(int i = head[u]; i != - 1; i = G[i].next)
        if(G[i].v != fa && ! vis[G[i].v])
            get_root(anc, G[i].v, u);
}
void get_dis(int u, int d, int fa)
{
    dis[num ++] = d;
    for(int i = head[u]; i != - 1; i = G[i].next)
        if(G[i].v != fa && ! vis[G[i].v])
            get_dis(G[i].v, d + G[i].w, u); 
}
int calc(int u, int d)
{
    int ret = 0;
    num = 0;
    get_dis(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)
{
    mn = n;
    get_size(u, - 1);
    get_root(u, u, - 1);
    ans += calc(root, 0);
    vis[root] = 1;
    for(int i = head[root]; i != - 1; i = G[i].next)
        if(! vis[G[i].v])
            ans -= calc(G[i].v, G[i].w), DFS(G[i].v);
}
int main()
{
    #ifndef ONLINE_JUDGE
    freopen("POJ1741.in", "r", stdin);
    freopen("POJ1741.out", "w", stdout);
    #endif
    int cas = 0;
    while((n = read()) && (k = read()))
    {
        cas ++;
        memset(head, - 1, sizeof(head));
        memset(vis, 0, sizeof(vis));
        ans = 0;
        for(int i = 1; i < n; i ++)
        {
            int u = read(), v = read(), w = read();
            add_edge(u, v, w);
            add_edge(v, u, w);
        }
        DFS(1);
        print(ans), putchar('
');
    }
}