P3806 【模板】点分治1

P3806 【模板】点分治1

给定一棵有 n 个点的树。

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

思路：点分治，先选取一个重心root，算出每个点到它的距离，从而可以通过枚举两条路径能够组成的 询问距离 的次数，但是这样会将同一子树上的两条路径算上(即路径长度加上了两倍子节点到root的距离)，将这种情况的贡献减去即可，这样就O(n)计算出了一颗经过root点的路径贡献，之后同样的对root的每一个相邻子树进行相同操作(总节点数仍是n)，因为每次选取的都是重心，因此只会有logn层，总复杂度O(nlogn)

#include<bits/stdc++.h>
#define MAXN 10010
#define INF 0x3f3f3f3f
using namespace std;
int head[MAXN],tot;
struct edge
{
    int v,w,nxt;
}edg[MAXN << 1];
inline void addedg(int u,int v,int w)
{
    edg[tot].v = v;
    edg[tot].w = w;
    edg[tot].nxt = head[u];
    head[u] = tot++;
}
int n,root,ms,sz[MAXN],Size;
bool vis[MAXN];
//root用于标记重心,ms表示树的重心的最大子树的大小
//sz[i]记录以i为根子树的大小,Size表示当前整棵树的大小,vis[i]表示当前节点是否被分治过
void getroot(int u,int f)//获得重心
{
    sz[u] = 1;
    int v,mson = 0;//mson记录以u为根最大子树的大小
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v] || v == f) continue;//剔除已经被分治过的点
        getroot(v,u);
        sz[u] += sz[v];
        if(sz[v] > mson) mson = sz[v];
    }
    if(Size - sz[u] > mson) mson = Size-sz[u];//把u看作根节点时u的父亲那一部分也算作子树
    if(ms > mson) ms = mson,root = u;//更新重心
}
int m,ask[110],ans[110];//所有询问
int dis[MAXN],cnt;//dis记录所有节点到重心的距离
struct node
{
    int d,num;
}nod[MAXN];//记录距离d出现的次数num
int cc;//不同距离数
void getdis(int u,int f,int d)//获得到目标点的距离
{
    dis[++cnt] = d;
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v] || v == f) continue;
        getdis(v,u,d + edg[i].w);
    }
}
void cal(int u,int d,int tp)//u表示getdis的起点,d表示u到目标点的距离,tp表示这一次统计出来的答案是合理的还是不合理的
{
    cnt = 0;
    getdis(u,0,d);//算出树中的点到目标点的距离
    sort(dis+1,dis+cnt+1);
    cc = 1,nod[cc].d = dis[1],nod[cc].num = 1;
    for(int i = 2;i <= cnt;++i)
        if(dis[i] != dis[i-1]) nod[++cc].d = dis[i],nod[cc].num = 1;
        else ++nod[cc].num;
    for(int i = 1;i <= m;++i)
    {
        int l = 1,r = cc;
        while(l <= r)
        {
            if(nod[l].d + nod[r].d == ask[i])
            {
                if(l == r) ans[i] += nod[l].num*(nod[l].num-1)/2 * tp;
                else ans[i] += nod[l].num * nod[r].num * tp;
                ++l,--r;
            }
            else if(nod[l].d + nod[r].d < ask[i])
                ++l;
            else
                --r;
        }
    }
}
void solve(int u,int ssize)//ssize是当前这棵子树的大小
{
    vis[u] = true;//代码保证每次进来的u都必定是当前这棵树的重心,我们将vis[u]标记为true,表示u点被分治过
    cal(u,0,1);//计算这棵树以u为重心的所有组合,但包括了共用同一条边的情况
    int v;
    for(int i = head[u];i != -1;i = edg[i].nxt)
    {
        v = edg[i].v;
        if(vis[v]) continue;
        cal(v,edg[i].w,-1);//将共用一条边的不合法情况去除
        ms = INF;//记得每次都要初始化
        Size = sz[v] < sz[u]?sz[v]:(ssize-sz[u]);//因为v实际上可能是u的父亲,故sz需相减
        //注意别把u,v写反了
        getroot(v,v);//求出以v为根节点的子树重心
        solve(root,Size);
    }
}
inline void init()
{
    tot = 0,ms = INF,Size = n;
    memset(head,-1,sizeof(int)*(n+1));
    memset(vis,false,sizeof(bool)*(n+1));
    memset(ans,0,sizeof(int)*(m+1));
}
int main()
{
    while(~scanf("%d%d",&n,&m))
    {
        init();
        int u,v,w;
        for(int i = 1;i < n;++i)
        {
            scanf("%d%d%d",&u,&v,&w);
            addedg(u,v,w),addedg(v,u,w);
        }
        for(int i = 1;i <= m;++i)
            scanf("%d",&ask[i]);
        getroot(1,1);
        solve(root,Size);
        for(int i = 1;i <= m;++i)
            if(ans[i] > 0) printf("AYE\n");
            else printf("NAY\n");
    }
    return 0;
}