CS Academy Max Score Tree

リンク : https://csacademy.com/contest/archive/task/max-score-tree/statement/

問題概要

N頂点の木とScoreiが与えられる。木のスコアは木に含まれる頂点vが持つ辺の数をEdge[v]とすると、この木のスコアはΣScore[Edge[v]]となる。入力で与えられた木の部分木が満たす最大のスコアを求めよ。

解説

dp[i]を自分の親を部分木として持つ場合の最大のスコアとする。持たない場合も別で計算しておくことも考えておく。頂点vに関して考えるとする。この頂点が持つ辺の集合をEとする。Eのそれぞれに関して、dp[e}(e∈E)をsortして上から貪欲にとって行くことを考える。上からi個足したものをsum[i]とすると、dp[v]=max{sum[i]+Score[i+1]}となる。もし親を持たないときはmax{sum[i]+Score[i]}となる。

LL n, score[110000], dp[110000], ans = -INF;
vector<vector<LL>>g,rvec;
LL func(int cur ,int par) {
    LL ret = score[1], sum = 0;
    if (dp[cur] != -INF)return dp[cur];
    vector<LL>d;
    rep(i, g[cur].size()) {
        LL v = g[cur][i];
        if (v == par)continue;
        func(v, cur);
        d.push_back(dp[v]);
    }
    sort(d.begin(), d.end(), [](LL x, LL y) {return x > y; });
    rep(i, d.size()) {
        sum += d[i];
        ret = max(ret, sum + score[i + 2]);
        ans = max(ans, sum + score[i + 1]);
    }
    return dp[cur] = ret;
}
int main() {
    cin >> n;
    rep(i, n)cin >> score[i];
    g.resize(n);
    rep(i, n - 1) {
        LL x, y;
        cin >> x >> y;
        x--, y--;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    rep(i, n)dp[i] = -INF;
    func(0, -1);
    cout << max(score[0],ans) << endl;
    return 0;
}