POJ - 1741 Tree

POJ - 1741 

思路：

点分治入门题

#include<iostream>
#include<cstdio>
#include<cstring>
#include<vector>
#include<algorithm>
#include<cmath>
#include<vector>
using namespace std;
#define fi first
#define se second
#define pi acos(-1.0)
#define LL long long
//#define mp make_pair
#define pb push_back
#define ls rt<<1, l, m
#define rs rt<<1|1, m+1, r
#define ULL unsigned LL
#define pll pair<LL, LL>
#define pii pair<int, int>
#define mem(a, b) memset(a, b, sizeof(a))
#define fio ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define fopen freopen("in.txt", "r", stdin);freopen("out.txt", "w", stout);
//head

const int N = 1e4 + 5;
int head[N], son[N], maxson[N], d[N], cnt = 0, ans = 0, rt, n, k;
bool vis[N];
vector<int>deep;
struct edge {
    int to, w, nxt;
}edge[N*2];
void add_edge(int u, int v, int w) {
    edge[cnt].to = v;
    edge[cnt].w = w;
    edge[cnt].nxt = head[u];
    head[u] = cnt++;
}
void get_rt(int o, int u) {
    son[u] = 1;
    maxson[u] = 0;
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        if(edge[i].to != o && !vis[edge[i].to]) {
            get_rt(u, edge[i].to);
            son[u] += son[edge[i].to];
            maxson[u] = max(maxson[u], son[edge[i].to]);
        }
    }
    maxson[u] = max(maxson[u], n - son[u]);
    if(maxson[u] < maxson[rt]) rt = u;
}
void get_d(int o, int u) {
    deep.pb(d[u]);
    son[u] = 1;
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        if(edge[i].to != o && !vis[edge[i].to]) {
            d[edge[i].to] = d[u] + edge[i].w;
            get_d(u, edge[i].to);
            son[u] += son[edge[i].to];
        }
    }
}
int cal(int u, int dis) {
    deep.clear();
    d[u] = dis;
    get_d(0, u);
    sort(deep.begin(), deep.end());
    int res = 0;
    for (int l = 0, r = deep.size()-1; l < r;) {
        if(deep[l] + deep[r] <= k) res += r - l++;
        else r--;
    }
    return res;
}
void work(int u) {
    ans += cal(u, 0);
    vis[u] = true;
    for (int i = head[u]; ~i; i = edge[i].nxt) {
        if(!vis[edge[i].to]) {

            ans -= cal(edge[i].to, edge[i].w);
            maxson[0] = n = son[edge[i].to];
            get_rt(rt = 0, edge[i].to);
            work(rt);
        }
    }
}
int main() {
    int u, v, w;
    while(~scanf("%d%d", &n, &k) && n && k) {
        mem(head, -1);
        mem(vis, false);
        ans = cnt = 0;
        for (int i = 1; i < n; i++) {
            scanf("%d%d%d", &u, &v, &w);
            add_edge(u, v, w);
            add_edge(v, u, w);
        }
        maxson[0] = n;
        get_rt(rt = 0, 1);
        work(rt);
        printf("%d\n", ans);
    }
    return 0;
}