BZOJ-2783: [JLOI2012]树

题目：http://www.lydsy.com/JudgeOnline/problem.php?id=2783

赤裸裸的水题，DFS一遍SET维护即可。

代码：

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <set>
 
using namespace std ;
 
#define ll long long
#define MAXN 100100
 
struct node {
    int k , cnt ;
    node(  ) {
        cnt = 0 ;
    }
    node( int _k , int _cnt ) : k( _k ) , cnt( _cnt ) {
         
    }
    bool operator < ( const node &a ) const {
        return k < a.k ;
    }
    bool operator == ( const node &a ) const {
        return k == a.k ;
    }
    bool operator > ( const node &a ) const {
        return k > a.k ;
    }
};
set < node > bst ;
 
struct edge {
    edge *next ;
    int t ;
} *head[ MAXN ] ;
 
void AddEdge( int s , int t ) {
    edge *p = new( edge ) ;
    p -> t = t , p -> next = head[ s ] ;
    head[ s ] = p ;
}
 
int n , len , w[ MAXN ] , roof ;
bool f[ MAXN ] ; 
ll ans = 0 ;
 
void dfs( int v , int sum ) {
    sum += w[ v ] ;
    set < node > :: iterator i = bst.find( node( sum - len , 0 ) ) ;
    if ( i != bst.end(  ) ) ans += ( ll )( i -> cnt ) ;
    i = bst.find( node( sum , 0 ) ) ;
    if ( i == bst.end(  ) ) bst.insert( node( sum , 1 ) ) ; else {
        node ret = *i ;
        bst.erase( i ) ;
        bst.insert( ret ) ;
    }
    for ( edge *p = head[ v ] ; p ; p = p -> next ) {
        dfs( p -> t , sum ) ;
    }
    i = bst.find( node( sum , 0 ) ) ;
    if ( i -> cnt == 1 ) bst.erase( i ) ; else {
        node ret = *i ;
        -- ret.cnt ;
        bst.erase( i ) ;
        bst.insert( ret ) ;
    }
}
 
int main(  ) {
    scanf( "%d%d" , &n , &len ) ;
    for ( int i = 0 ; i ++ < n ; ) scanf( "%d" , w + i ) ;
    memset( f , true , sizeof( f ) ) ;
    for ( int i = 1 ; i < n ; ++ i ) {
        int s , t ; scanf( "%d%d" , &s , &t ) ;
        AddEdge( s , t ) , f[ t ] = false ;
    }
    for ( int i = 0 ; i ++ < n ; ) if ( f[ i ] ) roof = i ;
    bst.clear(  ) ;
    bst.insert( node( 0 , 1 ) ) ;
    dfs( roof , 0 ) ;
    printf( "%lld\n" , ans ) ;
    return 0 ;
}