点分治模板题

阳光穿透心脏的1/2处 2023-10-11 11:48 51阅读 0赞

给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K。  
【题解】：  
　　我们找到树的重心，然后dfs，求出每个点到root的距离deep，然后对deep排序，扫描哪些点对是符合的。  
　　但是，点分治要求处理的路径是经过root，所以如果一条路径是在同一个子树之内的就不符合要求，所以还要对子树dfs一下，然后去重。  
　　接下来处理好root后，就可以处理其他连通块了，即递归其子树。

#include<cstdio>
    #include<cstring>
    #include<algorithm>
    using namespace std;
    
    #define N 10010
    #define inf 1e9+10
    struct node{int to,c,next;}g[N*2];
    int head[N],m;
    int son[N],f[N];
    bool vis[N];
    int d[N],deep[N];
    int n,sum,root,k,ans;
    
    void add_edge(int from,int to,int cost)
    {
        g[++m].next = head[from];
        head[from] = m;
        g[m].to = to; g[m].c = cost;
    }
    
    void getroot(int v,int fa)
    {
        son[v] = 1; f[v] = 0;
        for(int i = head[v];i;i=g[i].next)
            if(g[i].to != fa && !vis[g[i].to])
            {
                getroot(g[i].to,v);
                son[v] += son[g[i].to];
                f[v] = max(f[v],son[g[i].to]);
            }
        f[v] = max(f[v],sum - son[v]);
        if(f[v] < f[root]) root = v;
    }
    
    void getdeep(int v,int fa)
    {
        deep[++deep[0]] = d[v];
        for(int i = head[v];i;i=g[i].next)
            if(g[i].to != fa && !vis[g[i].to])
            {
                d[g[i].to] = d[v] + g[i].c;
                getdeep(g[i].to,v);
            }
    }
    
    int cal(int v,int cost)
    {
        d[v] = cost; deep[0] = 0;
        getdeep(v,0);
        sort(deep+1,deep+deep[0]+1);
        int l = 1,r = deep[0],sum = 0;
        while(l < r) {
            if(deep[l]+deep[r] <= k) {
                sum += r-l;
                l++;
            } else r--;
        }
        return sum;
    }
    
    void solve(int v)
    {
        ans += cal(v,0);
        vis[v] = 1;
        for(int i = head[v];i;i=g[i].next)
            if(!vis[g[i].to])
            {
                ans -= cal(g[i].to,g[i].c);
                sum = son[g[i].to];
                root = 0;
                getroot(g[i].to,0);
                solve(root);
            }
    }
    
    int main()
    {
        int u,v,w;
        while(scanf("%d%d",&n,&k) && n && k)
        {
            ans = root = m = 0;
            memset(vis,0,sizeof(vis));
            memset(head,0,sizeof(head));
            for(int i = 1;i < n;i++)
            {
                scanf("%d%d%d",&u,&v,&w);
                add_edge(u,v,w);
                add_edge(v,u,w);
            }
            f[0] = inf;
            sum = n;
            getroot(1,0);
            solve(root);
            printf("%d\n",ans);
        }
        return 0;
    }