POJ 1741 Tree

点分治

非常巧妙的实现, 复杂度 O(n * log ^2 n)
https://www.cnblogs.com/GXZlegend/p/6641720.html

#include <iostream>
#include <cstdio>
#include <cmath>
#include <algorithm>
#include <cstdlib>
#include <cstring>
using namespace std;
const int MAXN = 20005;
int head[MAXN], nume, n, k, siz[MAXN], masize, dep[MAXN];
int gra[MAXN], d[MAXN], tot, rot, ans;
bool f[MAXN];
int init() {
    int rv = 0, fh = 1;
    char c = getchar();
    while(c < '0' || c > '9') {
        if(c == '-') fh = -1;
        c = getchar();
    }
    while(c >= '0' && c <= '9') {
        rv = (rv<<1) + (rv<<3) + c - '0';
        c = getchar();
    }
    return fh * rv;
}
struct edge{
    int to, nxt, dis;
}e[MAXN];
void adde(int from, int to, int dis) {
    e[++nume].to = to;
    e[nume].nxt = head[from];
    e[nume].dis = dis;
    head[from] = nume;
}
void getroot(int u, int rt) {
    gra[u] = 0;
    siz[u] = 1;
    for(int i = head[u]; i; i = e[i].nxt) {
        int v = e[i].to;
        if(f[v] || v == rt) continue;
        getroot(v, u);
        siz[u] += siz[v];
        gra[u] = max(gra[u], siz[v]);
    }
    gra[u] = max(gra[u], masize - siz[u]);
    if(gra[rot] > gra[u]) rot = u;
}
void getdep(int u, int rt) {
    d[++tot] = dep[u];
    for(int i = head[u]; i; i = e[i].nxt) {
        int v = e[i].to;
        if(f[v] || v == rt) continue;
        dep[v] = dep[u] + e[i].dis;
        getdep(v, u);
    }
}
int calc(int u) {
    tot = 0;
    getdep(u, 0);
    sort(d + 1, d + 1 + tot);
    int l = 1, r = tot, rv = 0;
    while(l < r) {
        if(d[l] + d[r] <= k) rv += r - l, l++;
        else r--;
    }
    return rv;
}
void work(int u) {
    dep[u] = 0;
    f[u] = 1;
    ans += calc(u);
    for(int i = head[u]; i; i = e[i].nxt) {
        int v = e[i].to;
        if(!f[v]) {
            dep[v] = e[i].dis;
            ans -= calc(v);
            masize = siz[v];
            rot = 0;
            getroot(v, 0);
            work(rot);
        }
    }
}
int main() {
    while(1) {
        n = init(); k = init();
        if(!n || !k) break;
        memset(head, 0, sizeof(head));
        memset(f, 0, sizeof(f));
        nume = 0; tot = 0; masize = 0, ans = 0;
        for(int i = 1; i < n; i++) {
            int u = init(), v = init(), dis = init();
            adde(u, v, dis); adde(v, u, dis);
        }
        gra[0] = 0x3f3f3f3f;
        rot = 0;
        getroot(1, 0);
        work(rot);
        cout << ans << endl;
    }
    return 0;
}
posted @ 2018-04-01 16:44  Mr_Wolfram  阅读(...)  评论(...编辑  收藏