[BZOJ1316] 树上的询问

Description

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

Solution

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

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 10005;

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

void dfs1(int p,int from)
{
    siz[p]=1;
    msiz[p]=0;
    st.push_back(p);
    for(pair<int,int> 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<int,int> 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 <int> 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<int,int> 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<st.size(); i++)
    {
        msiz[st[i]]=max(msiz[st[i]],siz[p]-siz[st[i]]);
        if(msiz[st[i]]<rv) r=st[i],rv=msiz[st[i]];
    }
    return r;
}

void calc(int p)
{
    dis[p]=0;
    dfs2(p,0);
    buc.clear();
    buc.insert(0);
    for(pair<int,int> 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<int,int> 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<n; i++)
    {
        int t1,t2,t3;
        cin>>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")<<endl;
    }
}