洛谷P4178 Tree (点分治)

题目描述

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

输入输出格式

输入格式：

 

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

 

输出格式：

 

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

 

输入输出样例

输入样例#1： 

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

输出样例#1： 

5

题解：点分裸题，考虑分治中的暴力，将所有的重心子树中的点到中心的距离排序，对于一组l-r之间如果d[l]+d[r]<=k，显然d[l]+d[i](i<r)时都满足d[l]+d[i]<=k,可以统计答案，类似尺取的思想。
总复杂度是O(nlognlogn)

代码如下：

#include<map>
#include<set>
#include<queue>
#include<cmath>
#include<cstdio>
#include<string>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
#define mp make_pair
#define pii pair<int,int>
using namespace std;

vector<pii> g[];
int n,k,f[],vis[],size[],di[],cnt,ans;

void get_size(int now,int fa)
{
    size[now]=;
    f[now]=fa;
    for(int i=;i<g[now].size();i++)
    {
        if(vis[g[now][i].first]||g[now][i].first==fa) continue;
        get_size(g[now][i].first,now);
        size[now]+=size[g[now][i].first];
    }
}

int get_zx(int now,int fa)
{
    if(size[now]==) return now;
    int son,maxson=-;
    for(int i=;i<g[now].size();i++)
    {
        if(vis[g[now][i].first]||g[now][i].first==fa) continue;
        if(maxson<size[g[now][i].first])
        {
            maxson=size[g[now][i].first];
            son=g[now][i].first;
        }
    }
    int zx=get_zx(son,now);
    while(size[zx]<(size[now]-size[zx])*) zx=f[zx];
    return zx;
}

void get(int now,int fa,int dis)
{
    di[++cnt]=dis;
    for(int i=;i<g[now].size();i++)
    {
        if(vis[g[now][i].first]||g[now][i].first==fa) continue;
        get(g[now][i].first,now,dis+g[now][i].second);
    }
}

int calc(int now,int dis)
{
    cnt=;
    int tmp=;
    get(now,,dis);
    sort(di+,di+cnt+);
    int l=,r=cnt;
    while(l<=r)
    {
        if(di[l]+di[r]<=k)
        {
            tmp+=r-l;
            l++;
        }
        else r--;
    }
    return tmp;
}

void solve(int now)
{
    ans+=calc(now,);
    vis[now]=;
    for(int i=;i<g[now].size();i++)
    {
        if(vis[g[now][i].first]) continue;
        ans-=calc(g[now][i].first,g[now][i].second);
        get_size(g[now][i].first,);
        int zx=get_zx(g[now][i].first,);
        solve(zx);
    }
}

int main()
{
    scanf("%d",&n);
    for(int i=;i<n;i++)
    {
        int from,to,cost;
        scanf("%d%d%d",&from,&to,&cost);
        g[from].push_back(mp(to,cost));
        g[to].push_back(mp(from,cost));
    }
    scanf("%d",&k);
    solve();
    printf("%d\n",ans);
}