树形dp——没有上司的舞会

Description

没有上司的舞会 z5041

Sample Input 1

7
1
1
1
1
1
1
1
1 3
2 3
6 4
7 4
4 5
3 5

Sample Output 1

5

Algorithm

/*
树形DP
状态表示：
dp[u][0]表示在以u为根的子树中，选择不包含u结点的最大的快乐指数
dp[u][1]表示在以u为根的子树中，选择包含u结点的最大快乐指数
状态计算：
dp[u][0] = sum { max {dp[s][0], dp[s][1]} }, s为u的子结点
dp[u][1] = sum { dp[s][0] }, s为u的子结点
*/

Time Complexity

$O(n)$

Code Implementation

#include <iostream>
#include <cstring>

using namespace std;


const int N = 6010;

int dp[N][2], happy[N];

int h[N], ne[N], e[N], idx;
//fa[i]表示i结点父结点
int fa[N];

void add(int x, int y){
    e[idx] = y;
    ne[idx] = h[x];
    h[x] = idx;
    idx++;
}

void dfs(int u){
    //dp[u][1]表示在以u为根的子树中，选择包含u结点的最大快乐指数
    dp[u][1] = happy[u];
    //遍历u的所有子结点
    for(int i = h[u]; i != -1; i = ne[i]){
        int s = e[i];
        dfs(s);
        dp[u][0] += max(dp[s][0], dp[s][1]);
        dp[u][1] += dp[s][0];
    }
}

int main(){
    int n;

    scanf("%d", &n);

    for(int i = 1; i <= n; i++) scanf("%d", &happy[i]);

    //初始化树所有结点的链表头
    memset(h, -1, sizeof h);

    for(int i = 0; i < n - 1; i++){
        int a, b;
        scanf("%d%d", &a, &b);
        add(b, a);
        fa[a] = b; //b是a的父结点
    }

    int root = 1;
    while(fa[root]) root++; //找根结点



    dfs(root);

    printf("%d\n", max(dp[root][0], dp[root][1]));

    return 0;
}