洛谷 P3806 【模板】点分治1 点分治模版题

题目描述

给定一棵有n个点的树

询问树上距离为k的点对是否存在。

输入输出格式

输入格式：
n,m 接下来n-1条边a,b,c描述a到b有一条长度为c的路径

接下来m行每行询问一个K

输出格式：
对于每个K每行输出一个答案，存在输出“AYE”,否则输出”NAY”(不包含引号)

输入输出样例

输入样例#1：
2 1
1 2 2
2
输出样例#1：
AYE
说明

对于30%的数据n<=100

对于60%的数据n<=1000,m<=50

对于100%的数据n<=10000,m<=100,c<=1000,K<=10000000

分析：一道点分治的模版题。打得有点丑。

代码：

#include <iostream>
#include <cstdio>
#include <cmath>

const int maxn=10005;
const int inf=0x3f3f3f3;

using namespace std;

int n,m,root,cnt,x,y,w,t,head,tail,top,sum,i,j,pre;
int q[105],ans[105],ls[maxn],f[maxn],size[maxn],tree[maxn],vis[maxn],dis[maxn],g[maxn];
//q代表询问，ans代表询问答案，f表示最大子数大小，size表示子树大小，tree表示一个节点的队列，vis表示用过的根，dis表示距离，g表示子节点的栈
bool a[10000001];//一个桶

struct node{
    int y,w,next;
}edge[maxn*2];

void add_edge(int x,int y,int w)//加边
{
    edge[++cnt].y=y;
    edge[cnt].next=ls[x];
    edge[cnt].w=w;
    ls[x]=cnt;
    edge[++cnt].y=x;
    edge[cnt].next=ls[y];
    edge[cnt].w=w;
    ls[y]=cnt;
}

void getroot(int x,int fa)//寻找重心
{
    f[x]=0; size[x]=1;
    for (int t=ls[x];t>0;t=edge[t].next)
    {
        int y=edge[t].y;
        if ((y==fa) || (vis[y]==1)) continue;//保证不往回搜以及不搜用过的根
        getroot(y,x);
        size[x]+=size[y];//求以树上某一点得到的子树大小
        f[x]=max(f[x],size[y]);//最大子树大小
    }
    f[x]=max(f[x],sum-size[x]);
    if (f[x]<f[root]) root=x;//取最小的f作为新根
}

void dfs(int x,int fa)//做经过当前重心为根的链长
{
    size[x]=1;//以重心为根的子树大小
    for (int t=ls[x];t>0;t=edge[t].next)
    {
        int y=edge[t].y;
        if ((y==fa) || (vis[y]==1)) continue;
        dis[y]=dis[x]+edge[t].w;//更新每个点到重心的距离。
        g[++top]=dis[y];//把该距离压入栈
        dfs(y,x);
        size[x]+=size[y];//更新size
        if (x==root) //回到根
        {
        	for (int i=pre;i<=top;i++)
        	{
        		for (int j=1;j<=m;j++)
        		{
        			if (a[q[j]-g[i]])//判断之前的子树是否存在q[j]-g[i]的路径，有则更新答案
        			{
        				ans[j]=1;
        			}
        		}
        	}
        	for (int i=pre;i<=top;i++) a[g[i]]=true;//把当前子树的所有距离压入栈
        	pre=top+1;
        	tree[++tail]=y;//把该儿子作为下一棵树的根
        }
    }
}

int main()
{
    scanf("%d%d",&n,&m);
    for (i=1;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&w);
        add_edge(x,y,w);
    }
    for (i=1;i<=m;i++)
    {
    	scanf("%d",&q[i]);
    }      
    f[0]=inf; 
    head=0; tail=1;
    tree[tail]=1;//把1加入队列
    size[1]=n;       
   while (head<tail)
    {
        head++;
        sum=size[tree[head]];//sum为当前子树大小
        root=0;
        getroot(tree[head],0);//更新重心
        dis[root]=0;//重心距离改为0
        pre=1; top=0; a[0]=true;//把0改为可行距离    	  
        dfs(root,0);
        for (i=1;i<=top;i++) a[g[i]]=false;//把桶清空，其实就是把栈内数的对应位置清空
        vis[root]=1;//把根改为使用过
    }
    for (i=1;i<=m;i++)
    {
    	if (ans[i]) printf("AYE\n");
    	       else printf("NAY\n");
    }
}