【点分治】POJ 1741 Tree

第一次写点分治，一道入门题，稍微理解到了点分治解决的顺序和大致流程。

每次寻找当前子树的重心，围绕重心计算答案，这道题计算当前子树内经过了当前重心的满足条件的节点对数，用子树内总的符合条件的对数减去两个节点在同一子树内的对数。然后继续向子树分治。

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;

int bal, asize, sum, n, k, ans;

int tov[20004], nex[20004], h[10005], stot, w[20005];

void add ( int u, int v, int s ) {
    tov[++stot] = v;
    w[stot] = s;
    nex[stot] = h[u];
    h[u] = stot;
}

int siz[10005], vis[10005];

void find_root ( int u, int f ) {
    siz[u] = 1;
    int res = 0;
    for ( int i = h[u]; i; i = nex[i] ) {
        int v = tov[i];
        if ( v == f || vis[v] ) continue;必须要加vis条件，因为再往下分治时不能确定当前点的父亲是否被遍历过（点分治是围绕重心展开的！
        find_root ( v, u );
        siz[u] += siz[v];
        res = max ( res, siz[v] );
    }
    res = max ( res, sum - siz[u] );
    if ( res < asize ) {
        asize = res, bal = u;
    }
}

int dep[10004], dis[10005];

void get_dep ( int u, int f ) {
    dep[++dep[0]] = dis[u];
    siz[u] = 1;
    for ( int i = h[u]; i; i = nex[i] ) {
        int v = tov[i];
        if ( v == f || vis[v] ) continue;
        dis[v] = dis[u] + w[i];
        get_dep ( v, u );
        siz[u] += siz[v]; 
    }
}

int cal ( int u, int now ) {
    dis[u] = now; dep[0] = 0;
    get_dep ( u, 0 );
    sort ( dep + 1, dep + dep[0] + 1 );
    int tmp = 0, l = 1, r = dep[0];
    while ( l < r ) {
        if ( dep[l] + dep[r] <= k ) {
            tmp += r - l; l ++;///找能满足l的r统计对数
        } else r --;
    }
    return tmp;
}

void work ( int u ) {
    ans += cal ( u, 0 );
    vis[u] = 1;
    for ( int i = h[u]; i; i = nex[i] ) {
        int v = tov[i];
        if ( vis[v] ) continue;
        ans -= cal ( v, w[i] );
        sum = siz[v];
        asize = 0x3f3f3f3f;
        find_root ( v, u );
        work ( bal );
    }
}

int main ( ) {
    while ( scanf ( "%d%d", &n, &k ) == 2 ) {
        if ( n == 0 && k == 0 ) break;
        asize = 0x3f3f3f3f;
        stot = 0; ans = 0;
        memset ( h, 0, sizeof ( h ) );
        memset ( dis, 0, sizeof ( dis ) );
        memset ( vis, 0, sizeof ( vis ) );
        for ( int i = 1; i < n; i ++ ) {
            int a, b, c;
            scanf ( "%d%d%d", &a, &b, &c );
            add ( a, b, c );
            add ( b, a, c ); 
        }
        sum = n;
        find_root ( 1, 1 );
        work ( bal );
        printf ( "%d\n", ans );
    }
    return 0;
}