POJ 1741 Tree( Centroid D&C )

#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <set>
#include <utility>
#include <functional>
using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef pair< int, int > pii;
typedef vector< int > vi;
typedef vector< vi > vvi;
typedef vector< pii > vp;
typedef vector< vp > vvp;

const int INF = 0x3f3f3f3f;

// #define READ_FROM_FILE
// #define WRITE_TO_FILE
// #define MULTIPLE_CASES

void compute_subtree_size( int u, int fa, const vvp &G, vi &size, const vi &centroid ){
    size[ u ] = 1;
    for( int i = 0; i < G[ u ].size(); ++i ){
        int w = G[ u ][ i ].first;
        int v = G[ u ][ i ].second;
        if( v == fa || centroid[ v ] ) continue;
        compute_subtree_size( v, u, G, size, centroid );
        size[ u ] += size[ v ];
    }
}

pii get_centroid( int u, int fa, const vvp &G, const vi &subtree_size, const int root_size, vi &centroid ){
    pii res = pii( INF, 0 );
    int ch_max_size = 0, ch_sum_size = 0;
    for( int i = 0; i < G[ u ].size(); ++i ){
        int w = G[ u ][ i ].first;
        int v = G[ u ][ i ].second;
        if( v == fa || centroid[ v ] ) continue;
        res = min( res, get_centroid( v, u, G, subtree_size, root_size, centroid ) );
        ch_max_size = max( ch_max_size, subtree_size[ v ] );
        ch_sum_size += subtree_size[ v ];
    }
    res = min( res, pii( max( ch_max_size, root_size - ch_sum_size - 1 ), u ) );
    return res;
}

void enumerate_paths( int u, int fa, int len, const vvp &G, const vi &centroid, vi &dis ){
    dis.push_back( len );
    for( int i = 0; i < G[ u ].size(); ++i ){
        int w = G[ u ][ i ].first;
        int v = G[ u ][ i ].second;
        if( v == fa || centroid[ v ] ) continue;
        enumerate_paths( v, u, len + w, G, centroid, dis );
    }
}

int count_pairs( vi &dis, const int K ){
    int res = 0;
    sort( dis.begin(), dis.end() );
    for( int i = 0, j = dis.size(); i < dis.size(); ++i ){
        while( j > 0 && dis[ i ] + dis[ j - 1 ] > K ) --j;
        res += j - ( j > i );
    }
    return res / 2;
}

void sub_solve( int u, const vvp &G, int &ans, vi &centroid, const int K ){
    vi subtree_size( G.size() ); compute_subtree_size( u, -1, G, subtree_size, centroid );

    int c = get_centroid( u, -1, G, subtree_size, subtree_size[ u ], centroid ).second;
    centroid[ c ] = 1;

    for( int i = 0; i < G[ c ].size(); ++i ){
        int w = G[ c ][ i ].first;
        int v = G[ c ][ i ].second;
        if( centroid[ v ] ) continue;
        sub_solve( v, G, ans, centroid, K );
    }

    vi dis( 1, 0 ); // distance to c that is lower than K, first one is centroid itself
    for( int i = 0; i < G[ c ].size(); ++i ){
        int w = G[ c ][ i ].first;
        int v = G[ c ][ i ].second;
        if( centroid[ v ] ) continue;

        vi v_dis;
        enumerate_paths( v, c, w, G, centroid, v_dis );

        ans -= count_pairs( v_dis, K );
        dis.insert( dis.end(), v_dis.begin(), v_dis.end() );
    }

    ans += count_pairs( dis, K );
    centroid[ c ] = 0;
}

bool solve(){
    int N, K; cin >> N >> K;
    if( N == 0 and K == 0 ) return false;

    vvp G( N );
    for( int i = 0; i < N - 1; ++i ){
        int u, v, w; cin >> u >> v >> w;
        --u, --v;
        G[ u ].push_back( pii( w, v ) );
        G[ v ].push_back( pii( w, u ) );
    }

    int ans = 0;
    vi centroid( N );
    sub_solve( 0, G, ans, centroid, K );
    cout << ans << "\n";

    return true;
}

int main(){
    ios::sync_with_stdio( false );

    #ifdef READ_FROM_FILE
        freopen( "input.txt", "r", stdin );
    #endif

    #ifdef WRITE_FROM_FILE
        freopen( "output.txt", "w", stdout );
    #endif

    #ifdef MULTIPLE_CASES
        int T; cin >> T; while( T-- ) solve();
    #endif

    #ifndef MULTIPLE_CASES
        while( solve() )
            ;
    #endif

    return 0;
}