【HDU】5426 Rikka with Game【DP】

题目链接：【HDU】5426 Rikka with Game

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

typedef long long LL ;

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

const int MAXN = 100005 ;
const int MAXE = 200005 ;
const LL INF = 1e18 ;

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

Edge E[MAXE] ;
int H[MAXN] , cntE ;
int in[MAXN] , ou[MAXN] , dfs_idx ;
LL dp[MAXN][2] ;
int idx[MAXN] ;
LL val[MAXN] ;
int n , A , K ;

void init () {
    cntE = 0 ;
    dfs_idx = 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 ) {
    in[u] = ++ dfs_idx ;
    idx[dfs_idx] = u ;
    for ( int i = H[u] ; ~i ; i = E[i].n ) {
        int v = E[i].v ;
        if ( v == f ) continue ;
        dfs ( v , u ) ;
        val[u] += val[v] ;
    }
    ou[u] = dfs_idx ;
}

void solve () {
    int u , v ;
    init () ;
    for ( int i = 1 ; i <= n ; ++ i ) {
        scanf ( "%lld" , &val[i] ) ;
    }
    for ( int i = 1 ; i < n ; ++ i ) {
        scanf ( "%d%d" , &u , &v ) ;
        addedge ( u , v ) ;
        addedge ( v , u ) ;
    }
    dfs ( 1 , 0 ) ;
    LL ans = -INF ;
    dp[n + 1][0] = 0 ;
    dp[n + 1][1] = -INF ;
    for ( int k = 0 ; k <= min ( 1 , K ) ; ++ k ) {
        for ( int i = n , j = idx[i] ; i >= 2 ; j = idx[-- i] ) {
            dp[i][k] = max ( dp[i + 1][k] , dp[ou[j] + 1][k] - val[j] ) ;
            if ( k && dp[ou[j] + 1][k ^ 1] != -INF ) dp[i][k] = max ( dp[i][k] , dp[ou[j] + 1][k ^ 1] + ( A - 1 ) * val[j] ) ;
        }
        ans = max ( ans , dp[2][k] ) ;
    }
    dp[n + 1][0] = -INF ;
    for ( int k = 2 , cur = 0 ; k <= K ; ++ k , cur ^= 1 ) {
        for ( int i = n , j = idx[i] ; i >= 2 ; j = idx[-- i] ) {
            dp[i][cur] = max ( dp[i + 1][cur] , dp[ou[j] + 1][cur] - val[j] ) ;
            if ( dp[ou[j] + 1][cur ^ 1] != -INF ) dp[i][cur] = max ( dp[i][cur] , dp[ou[j] + 1][cur ^ 1] + ( A - 1 ) * val[j] ) ;
        }
        ans = max ( ans , dp[2][cur] ) ;
    }
    printf ( "%lld\n" , ans + val[1] ) ;
}

int main () {
    int T ;
    while ( ~scanf ( "%d%d%d" , &n , &A , &K ) ) solve () ;
    return 0 ;
}