[BZOJ1316] 树上的询问 - 点分治

Description

一棵 \(n\) 个点的带权有根树，有 \(p\) 个询问，每次询问树中是否存在一条长度为 \(len\) 的路径。\(n \le 10^4, p \le 10^2\)

Solution

点分治板子，用 set 维护，注意特判 \(k=0\) 的情况

#include 
using namespace std;
#define int long long
const int N = 10005;

vector > g[N];
// k : query
int dis[N],siz[N],msiz[N],n,m,k[N],ans[N],u[N];
vector  st;
vector  wl;

void dfs1(int p,int from)
{
    siz[p]=1;
    msiz[p]=0;
    st.push_back(p);
    for(pair pr:g[p])
    {
        int q=pr.first, w=pr.second;
        if(q!=from && u[q]==0)
        {
            dfs1(q,p);
            siz[p]+=siz[q];
            msiz[p]=max(msiz[p],siz[q]);
        }
    }
}

void dfs2(int p,int from)
{
    for(pair pr:g[p])
    {
        int q=pr.first, w=pr.second;
        if(q!=from && u[q]==0)
        {
            dis[q]=dis[p]+w;
            dfs2(q,p);
        }
    }
}

multiset  buc;

void dfs3(int p,int from)
{
    // Process each query
    for(int i=1; i<=m; i++)
    {
        if(buc.find(k[i]-dis[p]) != buc.end()) ans[i]=1;
    }
    for(pair pr:g[p])
    {
        int q=pr.first, w=pr.second;
        if(q!=from && u[q]==0)
        {
            dfs3(q,p);
        }
    }
    wl.push_back(p);
}

int findroot(int p)
{
    st.clear();
    dfs1(p,0);
    int r=0,rv=1e+9;
    for(int i=0; i pr:g[p])
    {
        // Process each subtree
        int q=pr.first, w=pr.second;
        if(u[q]==0)
        {
            dfs3(q,0);
            // wl: all nodes in this subtree
            for(int j:wl)
            {
                buc.insert(dis[j]);
            }
            wl.clear();
        }
    }
}

void solve(int p)
{
    int r=findroot(p);
    u[r]=1;
    calc(r);
    for(pair pr:g[r])
    {
        int q=pr.first, w=pr.second;
        if(u[q]==0)
        {
            solve(q);
        }
    }
}

signed main()
{
    ios::sync_with_stdio(false);
    cin>>n>>m;
    for(int i=1; i>t1>>t2>>t3;
        g[t1].push_back({t2,t3});
        g[t2].push_back({t1,t3});
    }
    for(int i=1; i<=m; i++) cin>>k[i];
    solve(1);
    for(int i=1; i<=m; i++)
    {
        cout<<(ans[i]||!k[i]?"Yes":"No")<