Tree POJ - 1741

Posted 2021-03-01 1024-xzx

题意：

给出一棵n个点的树，带权值，问两点之间简单路径长度 (leq k) 的点对的个数。

分析：

点分治。

代码【模板】：

#include <cstdio>
#include <algorithm>
#include <vector>
#include <cstring>
#include <string>
#define pb push_back
using namespace std;
typedef pair<int,int>P;
typedef long long ll;
const int N=1e4+5;
struct edge
{
    int to,nxt,len;
}G[N<<1];
int head[N],cot;
int dis[N],sz[N],focus,minn;
bool vis[N];//防止一个点重复找
int k,tn;//
ll ans;
void init()
{
    memset(head,-1,sizeof(head));
    cot=1;
}
void addedge(int from,int to,int len)
{
    G[cot].to=to;
    G[cot].len=len;
    G[cot].nxt=head[from];
    head[from]=cot++;
}
void read(int &x)
{
    x=0;
    int f=1;
    char ch=getchar();
    while(!isdigit(ch))
    {
        if(ch==‘-‘)
            f=-1;
        ch=getchar();
    }
    while(isdigit(ch))
    {
        x=(x<<3)+(x<<1)+ch-‘0‘;
        ch=getchar();
    }
    x*=f;
}
void dfs(int v,int p)//找重心v:根,p:父亲,
{
    sz[v]=1;
    int res=0;//最大子树的大小
    for(int i=head[v];i!=-1;i=G[i].nxt)
    {
        int u=G[i].to;
        if(u==p||vis[u])
            continue;
        dfs(u,v);
        sz[v]+=sz[u];
        res=max(res,sz[u]);
    }
    res=max(res,tn-sz[v]);//另外一棵子树
    if(res<minn)
    {
        minn=res;
        focus=v;
    }
}
void dfs2(int v,int p,int &cnt,int d)//求当前子树中各点离根的距离
{
    dis[++cnt]=d;
    for(int i=head[v];i!=-1;i=G[i].nxt)
    {
        int t=G[i].to;
        if(t==p||vis[t])
            continue;
        dfs2(t,v,cnt,d+G[i].len);
    }
}
int solve(int v,int d)//以重心为根进行处理
{
    if(k<2*d)
        return 0;
    int cnt=0,res=0;
    dfs2(v,v,cnt,0);//处理出子树中各点到根的距离
    sort(dis+1,dis+1+cnt);
    int l=1,r=cnt;
    while(l<=r)//类似二分，计数
    {
        if(dis[l]+dis[r]<=k-2*d)
        {
            res+=(r-l);
            l++;
        }
        else
            r--;
    }
    return res;
}
void divide(int v,int p)
{
    ans+=solve(v,0);
    vis[v]=1;
    for(int i=head[v];i!=-1;i=G[i].nxt)
    {
        int t=G[i].to;
        if(t==p||vis[t])
            continue;
        ans-=solve(t,G[i].len);
        minn=N;
        tn=sz[t];//注意子树的大小，是在该子树中找重心！！！！
        dfs(t,0);
        divide(focus,0);
    }
}
int main()
{
    int n,u,v,l;
    while(scanf("%d%d",&n,&k),n||k)
    {
        ans=0;
        init();
        memset(vis,false,sizeof(vis));
        for(int i=1;i<n;i++)
        {
            read(u),read(v),read(l);
            addedge(u,v,l);
            addedge(v,u,l);
        }
        tn=n,minn=N;
        dfs(1,0);//找到最开始树的重心
        divide(focus,0);
        printf("%lld
",ans);
    }
    return 0;
}