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\n",ans);
    return 0;
}