牛客小白月赛22 B 树上子链（树形dp）

 

示例1

输入

复制

5
2 -1 -1 -2 3
1 2
2 3
2 4
2 5

输出

复制

4

说明

样例中最大子链为1 -> 2 -> 5

备注:

一个结点，也可以称作一条链

好久没写树形dp了 场上这么简单的树形dp都没写出来

#include<bits/stdc++.h>
#define ll long long
const int maxn = 1e5 + 5;
int n;
int val[maxn];
int tot;
int head[maxn];
ll ans;
int vis[maxn];
int deg[maxn];
ll dp[maxn];
using namespace std;

struct node{
    int to;
    int next;
    node(){}
    node(int a, int b) : to(a), next(b) {}
}e[maxn << 1];

void init(){
    memset(head, -1, sizeof(head));
    tot = 0;
    memset(vis, 0, sizeof(vis));
}

void edgadd(int a, int b){
    e[tot] = node(b, head[a]);
    head[a] = tot++;
    e[tot] = node(a, head[b]);
    head[b] = tot++;
}

void dfs(int u, int pre){
    dp[u] = val[u];
    ll mxx = 0;
    ans = max(ans, dp[u]);
    for(int i = head[u]; i != -1; i = e[i].next){
        int v = e[i].to;
        if(v == pre)
            continue;
        dfs(v, u);
        ans = max(ans, mxx + dp[u] + dp[v]);
        mxx = max(mxx, dp[v]);
    }
    dp[u] += mxx;
}

int main(){
    init();
    cin >> n;
    for(int i = 1; i <= n; i++){
        cin >> val[i];
    }
    for(int i = 1;i <= n - 1; i++){
        int a, b;
        cin >> a >> b;
        edgadd(a, b);
        deg[a]++;
        deg[b]++;
    }
    int rt;
    int mx = -9999999;
    for(int i = 1; i <= n; i++){
        if(deg[i] == 1 && val[i] > mx){
            mx = val[i];
            rt = i;
        }
    }
    ans = mx;
    dfs(rt, 0);
    cout << ans << endl;
}