BZOJ1468Tree——点分治

题目描述

给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K

输入

N（n<=40000） 接下来n-1行边描述管道，按照题目中写的输入 接下来是k

输出

一行，有多少对点之间的距离小于等于k

样例输入

7
1 6 13
6 3 9
3 5 7
4 1 3
2 4 20
4 7 2
10

样例输出

5

点分治模板题，因为统计答案满足逆运算，所以可以用单步容斥来统计，将“任意两个子节点路径<=k的方案数”-“在同一子树内两节点路径<=k的方案数”就是最终答案。

#include<cmath>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
int n,k;
int tot;
int num;
int cnt;
int ans;
int root;
int x,y,z;
int s[80000];
int to[160000];
int mx[80000];
int val[160000];
int head[80000];
int next[160000];
int size[80000];
bool vis[80000];
void add(int x,int y,int v)
{
    tot++;
    next[tot]=head[x];
    head[x]=tot;
    to[tot]=y;
    val[tot]=v;
    tot++;
    next[tot]=head[y];
    head[y]=tot;
    to[tot]=x;
    val[tot]=v;
}
void getroot(int x,int fa)
{
    size[x]=1;
    mx[x]=0;
    for(int i=head[x];i;i=next[i])
    {
        if(to[i]!=fa&&!vis[to[i]])
        {
            getroot(to[i],x);
            size[x]+=size[to[i]];
            mx[x]=max(mx[x],size[to[i]]);
        }
    }
    mx[x]=max(mx[x],num-size[x]);
    if(!root||mx[x]<mx[root])
    {
        root=x;
    }
}
void dfs(int x,int fa,int dis)
{
    s[cnt++]=dis;
    for(int i=head[x];i;i=next[i])
    {
        if(to[i]!=fa&&!vis[to[i]])
        {
            dfs(to[i],x,dis+val[i]);
        }
    }
}
int calc(int x,int v)
{
    int ans=0;
    cnt=0;
    dfs(x,0,v);
    sort(s,s+cnt);
    for(int l=0,r=cnt-1;l<r;l++)
    {
        while(s[l]+s[r]>k&&l<r)
        {
            r--;
        }
        ans+=r-l;
    }
    return ans;
}
void partition(int x)
{
    vis[x]=1;
    ans+=calc(x,0);
    for(int i=head[x];i;i=next[i])
    {
        if(!vis[to[i]])
        {
            ans-=calc(to[i],val[i]);
            num=size[to[i]];
            root=0;
            getroot(to[i],0);
            partition(root);
        }
    }
}
int main()
{
    mx[0]=2147483647;
    scanf("%d",&n);
    for(int i=1;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);
    }
    scanf("%d",&k);
    getroot(1,0);
    partition(root);
    printf("%d",ans);
}