P3806 【模板】点分治1

题目背景

感谢hzwer的点分治互测。

题目描述

给定一棵有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

//Pro:P3806 【模板】点分治1

//写一篇与他们做法不同的题解
//这个做法像 P4178 Tree 差不多 

//仍然是存下所有的询问，离线操作
//我们用flag[i]表示第i个询问能满足的次数
//在分治的时候，仍然是先处理父亲，再容斥处理减掉儿子 

//如果都处理完之后flag[i]仍然是>0的，那就说明存在这样一条路径，否则就不存在

//复杂度的话，应该是m*nlogn+nlog^2n的吧 不太会算 

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;

inline int read()
{
    char c=getchar();int num=0;
    for(;!isdigit(c);c=getchar());
    for(;isdigit(c);c=getchar())
        num=num*10+c-'0';
    return num;
}

const int N=1e4+5;
const int M=1e7+5;
const int INF=599518803;

int n,m;
int k[N];

int head[N],num_edge;
struct Edge
{
    int v,w,nxt;
}edge[N<<1];

inline void add_edge(int u,int v,int w)
{
    edge[++num_edge].v=v;
    edge[num_edge].w=w;
    edge[num_edge].nxt=head[u];
    head[u]=num_edge;
}

bool vis[N];
int flag[M];

int Siz,root;
int mxson[N],siz[N];
void getroot(int u,int fa)
{
    siz[u]=1,mxson[u]=0;
    for(int i=head[u],v;i;i=edge[i].nxt)
    {
        v=edge[i].v;
        if(v==fa||vis[v])
            continue;
        getroot(v,u);
        siz[u]+=siz[v];
        mxson[u]=max(mxson[u],siz[v]);
    }
    mxson[u]=max(mxson[u],Siz-siz[u]);
    if(mxson[u]<mxson[root])
        root=u;
}

int dep[N],dis[N],cnt;
void getdis(int u,int fa)
{
    dis[++cnt]=dep[u];
    for(int i=head[u],v;i;i=edge[i].nxt)
    {
        v=edge[i].v;
        if(vis[v]||v==fa)
            continue;
        dep[v]=dep[u]+edge[i].w;
        getdis(v,u);
    }
}

void solve(int u,int dist,int val)      //val表示是加还是容斥减掉 
{
    dep[u]=dist,cnt=0;
    getdis(u,u);
    sort(dis+1,dis+cnt+1);
    for(int i=1,l,r;i<=m;++i)
    {
        for(l=1,r=cnt;l<r;)
        {
            if(dis[l]+dis[r]>k[i])      //调整边界 
                --r;
            else if(dis[l]+dis[r]<k[i])
                ++l;
            else    //可以拼出了，计算点对个数 
            {
                if(dis[l]==dis[r])
                {
                    flag[i]+=val*(r-l+1)*(r-l)/2;
                    break;
                }
                else
                {
                    int l1=1,l2=1;      //枚举求相等的数有多少个 
                    for(;l<r&&dis[l+1]==dis[l];++l,++l1);
                    for(;r>l&&dis[r-1]==dis[r];--r,++l2);
                    flag[i]+=val*l1*l2;
                    ++l,--r;
//                    int ll=l,rr=r,mid,pos=ll;     二分求 
//                    while(ll<=rr)
//                    {
//                        mid=(ll+rr)>>1;
//                        if(dis[mid]==dis[l])
//                        {
//                            ll=mid+1;
//                            pos=mid;
//                        }
//                        else
//                            rr=mid-1;
//                    }
//                    int a=pos-l+1;
//                    l=pos+1;
//                    ll=l,rr=r,pos=rr;
//                    while(ll<=rr)
//                    {
//                        mid=(ll+rr)>>1;
//                        if(dis[mid]==dis[r])
//                            rr=mid-1,pos=mid;
//                        else
//                            ll=mid+1;
//                    }
//                    flag[i]+=val*a*(r-pos+1);
//                    r=pos-1;
                }
            }
        }
    }
}

void divide(int u)
{
    solve(u,0,1);   //算重心的答案 
    vis[u]=1;
    for(int i=head[u],v;i;i=edge[i].nxt)
    {
        v=edge[i].v;
        if(vis[v])
            continue;
        solve(v,edge[i].w,-1);      //将儿子减掉 
        Siz=siz[v],root=0;
        getroot(v,v);
        divide(root);
    }
}

int main()
{
    n=read(),m=read();
    for(int i=1,a,b,c;i<n;++i)
    {
        a=read(),b=read(),c=read();
        add_edge(a,b,c);
        add_edge(b,a,c);
    }
    for(int i=1;i<=m;++i)
        k[i]=read();
    Siz=n,mxson[0]=INF;
    getroot(1,1);
    divide(root);
    for(int i=1;i<=m;++i)
    {
        if(flag[i])
            puts("AYE");
        else
            puts("NAY");
    }
    return 0;
}