【点分治】POJ 1741 Tree

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

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

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

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

int tov[], nex[], h[], stot, w[];

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

int siz[], vis[];

void find_root ( int u, int f ) {
    siz[u] = ;
    int res = ;
    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[], dis[];

void get_dep ( int u, int f ) {
    dep[++dep[]] = dis[u];
    siz[u] = ;
    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[] = ;
    get_dep ( u,  );
    sort ( dep + , dep + dep[] +  );
    int tmp = , l = , r = dep[];
    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,  );
    vis[u] = ;
    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 = ;
        find_root ( v, u );
        work ( bal );
    }
}

int main ( ) {
    while ( scanf ( "%d%d", &n, &k ) ==  ) {
        if ( n ==  && k ==  ) break;
        asize = ;
        stot = ; ans = ;
        memset ( h, , sizeof ( h ) );
        memset ( dis, , sizeof ( dis ) );
        memset ( vis, , sizeof ( vis ) );
        for ( int i = ; 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 ( ,  );
        work ( bal );
        printf ( "%d\n", ans );
    }
    return ;
}