luogu P4178 Tree

题目链接

luogu P4178 Tree

题解

点分治

代码

// luogu-judger-enable-o2
#include<cstdio> 
#include<algorithm> 
inline int read() { 
    int x = 0,f = 1 ; 
    char c = getchar(); 
    while(c < '0' || c > '9') c = getchar(); 
    while(c <= '9' && c >= '0') x = x * 10 + c - '0',c = getchar(); 
    return x * f; 
} 
#define LL long long
const int maxn = 80007; 
struct node {
    int v,w,next; 
} edge[maxn << 1]; 
int head[maxn],num = 0; 
inline void add_edge(int u,int v,int w) { 
    edge[++ num].v = v; edge[num].next = head[u];edge[num]. w= w; head[u] = num; 
} 
int n,K;
int root = 0,f[maxn],siz[maxn],g[maxn];  
bool vis[maxn]; 
int tot; 
void getroot(int x,int fa) { 
    siz[x] = 1; int num = 0; f[x] = 0; 
    for(int i = head[x];i;i = edge[i].next) { 
 		int v = edge[i].v; 
        if(v == fa || vis[v]) continue; 
        getroot(v,x); 
        siz[x] += siz[v]; f[x] = std::max(f[x],siz[v]); 
    } 
    f[x] = std::max(f[x],tot - siz[x]);  
    if(f[root] > f[x]) root = x;  
} 
int l,r; 
int nod[maxn],dis[maxn]; 
void getdis(int x,int fa) { 
    nod[++ r] = dis[x]; 
    for(int i = head[x];i;i = edge[i].next) { 
        int v = edge[i].v; 
        if(v == fa || vis[v]) continue; 
        dis[v] = dis[x] + edge[i].w; 
        getdis(v,x); 
    } 
} 
LL calc(int x,int y) { 
    dis[x] = y; 
    r = 0,l = 1; 
    getdis(x,0); 
    LL ret = 0; 
    std::sort(nod + 1,nod + r + 1); 
    for(;l < r;) 
        if(nod[l] + nod[r] <= K) ret += r - l,l ++; 
        else r --; 
    return ret; 
} 
LL ans = 0; 
void solve(int x) { 
    ans += calc(x,0); 
    vis[x] = true; 
    for(int i = head[x];i; i = edge[i].next) { 
        int v = edge[i].v; 
        if(vis[v]) continue; 
        ans -= calc(v,edge[i].w); 
        tot = siz[v]; root = 0; 
        getroot(v,0); 
        solve(root); 
    } 
} 
int main() { 
    f[0] = 0x3f3f3f3f; 
    n = read(); 
    for(int u,v,w,i = 1;i < n;++ i) { 
        u = read(),v = read(),w = read(); 
        add_edge(u,v,w); add_edge(v,u,w); 
    }  
    K = read(); 
    tot = n; 
    getroot(1,0); 
    solve(root); 
    printf("%lld
",ans); 
    return 0; 
}