【HDU】5905 Black White Tree【树dp，枚举子树，对于一个a，其对应的可行b一定是一个连续区间】

题目链接：Black White Tree

枚举子树，对于一个a，其对应的可行b一定是一个连续区间。

#include <bits/stdc++.h>
using namespace std ;

typedef long long LL ;
typedef pair < int , int > pii ;

#define clr( a , x ) memset ( a , x , sizeof a )

const int MAXN = 2005 ;
const int MAXE = 4005 ;
const int INF = 0x3f3f3f3f ;

struct Edge {
    int v , n ;
    Edge () {}
    Edge ( int v , int n ) : v ( v ) , n ( n ) {}
} ;

Edge E[MAXE] ;
int H[MAXN] , cntE ;
pii G[MAXN][MAXN] , nxt[MAXN] ;
int n ;
int val[MAXN] , b , w ;
int c[MAXN][MAXN] ;
char s[MAXN] ;
int wa[MAXN] ;

void init () {
    cntE = 0 ;
    clr ( H , -1 ) ;
}

void addedge ( int u , int v ) {
    E[cntE] = Edge ( v , H[u] ) ;
    H[u] = cntE ++ ;
}

void dfs ( int u , int f ) {
    wa[u] = val[u] == 0 ;
    for ( int i = H[u] ; ~i ; i = E[i].n ) {
        int v = E[i].v ;
        if ( v == f ) continue ;
        dfs ( v , u ) ;
        wa[u] += wa[v] ;
    }
    int L = val[u] == 0 , R = val[u] == 0 ;
    G[u][L] = pii ( val[u] == 1 , val[u] == 1 ) ;
    G[u][L ^ 1] = pii ( INF , 0 ) ;
    for ( int i = H[u] ; ~i ; i = E[i].n ) {
        int v = E[i].v ;
        if ( v == f ) continue ;
        for ( int j = 0 ; j <= R + wa[v] ; ++ j ) {
            nxt[j] = pii ( INF , 0 ) ;
        }
        for ( int j = L ; j <= R ; ++ j ) {
            int l = G[u][j].first , r = G[u][j].second ;
            for ( int k = 0 ; k <= wa[v] ; ++ k ) {
                int x = G[v][k].first , y = G[v][k].second ;
                int nl = l + x , nr = r + y ;
                int X = min ( nl , nxt[j + k].first ) ;
                int Y = max ( nr , nxt[j + k].second ) ;
                nxt[j + k] = pii ( X , Y ) ;
            }
        }
        R += wa[v] ;
        for ( int j = 0 ; j <= R ; ++ j ) {
            G[u][j] = nxt[j] ;
        }
    }
    for ( int i = L ; i <= R ; ++ i ) {
        int l = G[u][i].first , r = G[u][i].second ;
        c[i][l] ++ ;
        c[i][r + 1] -- ;
    }
    G[u][0] = pii ( 0 , G[u][0].second ) ;
    c[0][0] ++ ;
    c[0][1] -- ;
}

void solve () {
    b = w = 0 ;
    init () ;
    scanf ( "%d" , &n ) ;
    scanf ( "%s" , s ) ;
    for ( int i = 0 ; i <= n ; ++ i ) {
        for ( int j = 0 ; j <= n ; ++ j ) {
            c[i][j] = 0 ;
        }
    }
    for ( int i = 0 ; i < n ; ++ i ) {
        if ( s[i] == '0' ) w ++ ;
        else b ++ ;
        val[i + 1] = s[i] - '0' ;
    }
    for ( int i = 1 ; i < n ; ++ i ) {
        int u , v ;
        scanf ( "%d%d" , &u , &v ) ;
        addedge ( u , v ) ;
        addedge ( v , u ) ;
    }
    dfs ( 1 , 1 ) ;
    LL ans = 0 ;
    for ( int i = 0 ; i <= n ; ++ i ) {
        ans += ( c[i][0] > 0 ) * ( i + 1 ) ;
        for ( int j = 1 ; j <= n ; ++ j ) {
            c[i][j] += c[i][j - 1] ;
            ans += ( c[i][j] > 0 ) * ( i + 1 ) * ( j + 1 ) ;
        }
    }
    printf ( "%lld\n" , ans ) ;
}

int main () {
    int T ;
    scanf ( "%d" , &T ) ;
    for ( int i = 1 ; i <= T ; ++ i ) {
        solve () ;
    }
    return 0 ;
}