cqyz oj | 树的分治 | 树形DP | 树的重心

Posted 2022-05-19 de-compass

Description

　　给定一棵N个节点的带权树，定义dist(u,v)为u,v两点间的最短路径长度，路径的长度义为路径上所有边的权和。再给定一个K，如果对于不同的两个结点a,b，如果满足dist(a,b)<=K，则称(a,b)为合法点对。
　　你的任务是求合法点对个数。

Input

　　第一行包含两个个整数N和K，接下来的N-1行，每行包含三个整数：u,v,len，表示树边(u,v)的长度len。

Output

　　一个整数，表示合法点对的数目。

Sample Input 1

5 4
1 2 3
1 3 1
1 4 2
3 5 1

Sample Output 1

8

Hint

0<N<100 000,0<K<1000000000,0<len<10000

---

看到数据范围先想到用long long
直接跳过O(n^2)暴力算法

满足条件的点a和b有以下两种情况
一、a,b之间的路径经过根，即a,b在根节点的两棵不同子树里，或者a,b其中之一是根节点
1.计算root到各点距离dis[j]
2.在dis[]中统计满足dis[a]+dis[b]<=K的点对数量X，转化为化装晚会一题
3.去掉来自同一子树中的点对数Y
答案为X-Y
二、a,b在之间的路径不经过root，即a,b在同一棵子树里
这种情况就是分治算法的子问题，可以递归求
（为什么要先X-Y再在子树里重新求呢？因为由“一、”里计算出来的点对本来可以不经过root，相当于多算了一段长度，导致有些原本满足条件的点对因为这截多算的长度而未被统计到）

由这个思路，设如下函数：
getdist(i)计算以i为根的子树中i到各点的距离 \\(O(size[i])\\)
F(i,k)计算以i为根的子树中dis[a]+dis[b]<=k的点对数量\\(O(size[i]log_2size[i])\\)
calc(i)递归函数 \\(O(n)\\)
该方法时间复杂度\\(O(depth * nlogn)\\)，如果随便找一个根可能导致depth为n，时间退化为\\(O(n^2logn)\\)

要使depth最小，想到找树的重心，因此又设一个函数getroot(i)找i所在的子树里的重心为根，\\(O(size[i])\\)
这样时间稳定在\\(O(n*(log_2n)^2)\\)

此题代码实现对于我这种蒟蒻有一定难度，有很多函数调用之间的细节：

#include<cstdio>
#include<iostream>
#include<queue>
#include<cstring>
#include<algorithm>
#include<cmath>
#define maxn 100005 
using namespace std;
typedef long long ll;
int fir[maxn],ne[maxn*2],to[maxn*2],w[maxn*2],newp=0;
void add(int x,int y,int z)
    ne[++newp]=fir[x];
    fir[x]=newp;
    to[newp]=y;
    w[newp]=z;

int siz[maxn],mn,rt;
bool mark[maxn]=0;
void getroot(int u,int f,int tot)
    siz[u]=1;
    int tmp=0;
    for(int i=fir[u];i;i=ne[i])
        int v=to[i];
        if(v==f || mark[v])continue;
        getroot(v,u,tot);
        siz[u]+=siz[v];
        tmp=max(tmp,siz[v]);
    
    tmp=max(tmp,tot-siz[u]);
    if(tmp<mn)rt=u,mn=tmp;


ll dis[maxn];
int cnt;
void getdist(int u,int f,ll dist)
    dis[cnt++]=dist;
    for(int i=fir[u];i;i=ne[i])
        int v=to[i];
        if(v==f || mark[v])continue;
        getdist(v,u,dist+w[i]);
    

ll F(int u,int k)
    cnt=0;
    getdist(u,0,0);
    sort(dis,dis+cnt);
    ll res=0;
    for(int i=0;i<cnt;i++)
        res+=upper_bound(dis,dis+i,k-dis[i])-dis;
    return res;

int n,K;
ll calc(int u)
    ll t=F(u,K);
    mark[u]=1;
    for(int i=fir[u];i;i=ne[i])
        int v=to[i];
        if(mark[v])continue;
        t-=F(v,K-2*w[i]);
        mn=n;
        getroot(v,0,siz[v]);
        t+=calc(rt);
    
    return t;

int main()
    scanf("%d%d",&n,&K);mn=n;
    for(int i=1,x,y,z;i<n;i++)
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z);add(y,x,z);
    
    getroot(1,0,n);
    printf("%lld",calc(rt));
    return 0;