P4178 Tree

第一道点分治的题目,感觉代码难写hhh,大概思路就是染色法+双指针模拟+分治吧~//..

#include <bits/stdc++.h>
using namespace std;
const int N=4e4+40;
struct node{
    int to,val;
};
vector<node>G[N];
int root,color[N],f[N],dp[N],sum,vis[N],ans,siz[N],dis[N],heap[N],id,k,n;
bool cmp(int a,int b)
{
    return dis[a]<dis[b];
}

void f_dis(int u,int pre)
{
    //     puts("1");
    heap[++id]=u;
    color[u]=color[pre];f[color[u]]++;
    for(int i=0;i<G[u].size();i++)
    {
        int v=G[u][i].to;
        if(vis[v]||v==pre)  continue;
        dis[v]=G[u][i].val+dis[u];
        f_dis(v,u);
    }
}

int cal(int u)//计算以这个根的分治数.
{
    //puts("1");
    dis[u]=0;id=0;heap[++id]=u;
    color[u]=u;f[color[u]]=1;
    for(int i=0;i<(int)G[u].size();i++)
    {
        //puts("1");
        int v=G[u][i].to;
        if(vis[v])  continue;
        dis[v]=G[u][i].val;color[v]=v;f[color[v]]=0;
        f_dis(v,v);
    }//puts("1");
    sort(heap+1,heap+1+id,cmp);
    int res=0,l=1,r=id;
    while(l<r)
    {
        while(l<r&&dis[heap[r]]+dis[heap[l]]<=k)
        {
            res+=r-l-f[color[heap[l]]]+1;
            f[color[heap[l]]]--;
            l++;
        }
        f[color[heap[r]]]--;
        r--;
    }
    return res;
}

void f_root(int u,int pre)
{
    siz[u]=1;dp[u]=0;
    for(int i=0;i<(int)G[u].size();i++)
    {
        int v=G[u][i].to;
        if(v==pre||vis[v])  continue;
        f_root(v,u);siz[u]+=siz[v];
        dp[u]=max(dp[u],siz[v]);
    }dp[u]=max(dp[u],sum-siz[u]);
    if(dp[u]<dp[root])  root=u;
}

void solve(int u)
{
    vis[u]=1;
    ans+=cal(u);
    for(int i=0;i<(int)G[u].size();i++)
    {
        int v=G[u][i].to;
        if(vis[v])  continue;
        sum=siz[v];dp[0]=n,root=0;
        f_root(v,root);
        solve(root);
    }
}

int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n-1;i++)
    {
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        G[u].push_back({v,w});
        G[v].push_back({u,w});
    }dp[0]=sum=n;scanf("%d",&k);
    f_root(1,0);
    solve(root);
    printf("%d\n",ans);
    return 0;
}