题目: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 ;
}