点分治学习笔记

· · 个人记录

概念

点分治用于解决有一定要求的链的计数。

对于点 u 的子树的问题,可以将答案分为:

  1. 经过点 u

  2. 不经过点 u

第一种可以用桶加暴力。枚举一端的长度,用桶计算另一端长度;第二种分到子树中解决即可。

注意到,在随机选根的时候该算法表现不优秀,但若根为重心,因为每次子树大小都减少一半,所以时间复杂度优秀,为 O(n\log n)

代码

#include <iostream>
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <vector>
#include <algorithm>
#include<set>
#include<queue>
#define mod 998244353
#define int long long
using namespace std;
int n,m,ans[10001],ques[10001],use[10001],len,fin[10001],len2,ton[10000001],anss,mid,f[10001];
bool debug=false;
struct edge
{
    int v,w;
};
vector<edge> g[10001];
bool vis[10001];
void getmid(int u,int fa,int sz) //重心 
{
    int ma=0;
    for(int i=0;i<g[u].size();i++)
    {
        if(g[u][i].v==fa||vis[g[u][i].v]) continue;
        getmid(g[u][i].v,u,sz);
        f[u]+=f[g[u][i].v]+1; //更新状态
        ma=max(ma,f[g[u][i].v]); //最大子节点 
    }
    ma=max(ma+1,sz-f[u]-1);
    if(ma<anss) //如果小于当前 
    {
        anss=ma;
        mid=u;
    }
}
void init()
{
    for(int i=1;i<=len;i++)
    {
        ton[use[i]]=0; //清空桶
    }
    len=0;
}
void dfs(int u,int fa,int tot) //找值
{
    fin[++len2]=tot;
    for(int i=0;i<g[u].size();i++)
    {
        if(g[u][i].v==fa||vis[g[u][i].v]) continue;
        dfs(g[u][i].v,u,tot+g[u][i].w);
    }
}
int getsiz(int u,int fa) //因为换根,子树大小要重新计算
{
    int ret=1;
    f[u]=0;
    for(int i=0;i<g[u].size();i++)
    {
        if(g[u][i].v==fa||vis[g[u][i].v]) continue;
        ret+=getsiz(g[u][i].v,u);
    }
    return ret;
}
void work(int u)
{
    for(int i=0;i<g[u].size();i++)
    {
        if(vis[g[u][i].v]) continue;
        len2=0;
        dfs(g[u][i].v,u,g[u][i].w);
        for(int z=1;z<=len2;z++)
        {
            for(int j=1;j<=m;j++)
            {
                if(ques[j]-fin[z]>=0&&ques[j]-fin[z]<=10000000) ans[j]+=ton[ques[j]-fin[z]]; //有,计算一下
            }
        }
        for(int z=1;z<=len2;z++)
        {
            if(fin[z]>10000000) continue;
            ton[fin[z]]++;
            if(ton[fin[z]]==1) use[++len]=fin[z];
        }
    }
    for(int j=1;j<=m;j++)
    {
        ans[j]+=ton[ques[j]]; //单独的链
    }
}
void dfz(int u)
{
    anss=10000000000;
    mid=0;
    getmid(u,-1,getsiz(u,-1)); //重心
    init();
    u=mid; //换根
    work(u);
    vis[u]=true;
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i].v;
        if(vis[v]) continue;
        dfz(v);
    }
}
signed main()
{
    scanf("%lld%lld",&n,&m);
    for(int i=1,u,v,w;i<n;i++)
    {
        scanf("%lld%lld%lld",&u,&v,&w);
        g[u].push_back((edge){v,w});
        g[v].push_back((edge){u,w});
    }
    for(int i=1;i<=m;i++)
    {
        scanf("%lld",&ques[i]);
    }
    dfz(1);
    for(int i=1;i<=m;i++)
    {
        cout<<(ans[i]==0?"NAY\n":"AYE\n");
    }
}