743-D. Chloe and pleasant prizes

2021-03-19 12:21:41 +0800 CST Tags: dp, tree

出處：https://codeforces.com/contest/743/problem/D
提交：https://codeforces.com/contest/743/submission/110412498
提交：https://codeforces.com/contest/743/submission/111250285 (沒有把答案用 DP 存起來的版本)

問題：給定 $n\in[1,2\cdot10^5]$ 個禮物以及它們各自介於 $[-10^9,10^9]$ 的受喜愛程度。接下來的 n-1 行說明禮物 u 和禮物 v 會被一條繩子牽連在一起 (注意沒有指明誰上誰下)，最後保證這整棵聖誕樹會以禮物 1 為根結點。現在兩人要同時從這棵樹上挑選禮物，挑選的方式是只要選定一個禮物，那麼所有掛在它底下 (延伸到葉子節點) 的禮物也要通通一起選出來，兩人的禮物當然不可重複 (兩棵子樹不可有交點)，否則會吵架，請問如何調配能使得所有入選禮物的受喜愛程度總和最大？最大值為何？

解法：
這題的解法很經典地會使用 DP，而且不會太過困難，其實只要 1-pass DFS 就可以直接把答案記錄在頂點中。而當初卡住我思路的點的是在於如何同時維護這棵子樹底下 (一棵有最大和的子樹值) 以及 (兩棵子樹的最大和) 這兩種資訊，事實上我們只要同時維護三個陣列 (sum: 單純子樹底下的禮物喜愛度總和)、(one: 子樹底下的一棵最大和子樹值)、(two: 子樹底下的兩棵子樹的最大和) 即可。第一個 sum 陣列很簡單，如果是葉子的話直接存入自己禮物值，不是的話就先把直屬小孩 (children) 的 sum 全部加起來，再和自己禮物作加總並存入自己的 sum 值；第二個 one 陣列稍微要想一下，我們可以分成選或不選自己這個禮物，如果不選的話就是從 children 的 one 裡面挑最大的那一個，如果要選的話就是要再檢查自己的 sum 會不會更大，有更大就選，否則就不能選；而第三個 two 陣列是最不直覺的，我可能是某一個 child 底下就已經選好兩棵樹了，也有可能這個 child 底下的一棵最大子樹和另外一個 child 底下的一棵最大子樹才能拼成此時的最大和，由此觀察，我們其實可以先找哪個 child 有最大的 two 陣列值，然後「如果至少有 2 個 children 的話」再和所有 children 裡面有最大 one 和次大 one 的總和去作比較，得到這個節點真正的 two 值。最後根節點的 two 值就是我們要的。

另外，有讀者可能會發現因為每個子樹的答案只會用到一遍，就是 DFS 回傳答案的那一刻，所以根本毋須再另外開三個 DP 陣列 sum、one、two，我有再丟一次不用 DP 的版本，但是不知道為什麼記憶體用量不減反增，連結在上面，有興趣的讀者可以自行參考。

延伸：如果所求子樹不用延伸到葉子節點，可以在中間就腰斬，其他條件不變，那麼程式又該如何撰寫？

實作：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44


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

int n, u, v, a[200001];
long long sum[200001], one[200001], two[200001];
vector<int> adj[200001];

void dfs(int u) {
    long long one_firstmax = LLONG_MIN, one_secondmax = LLONG_MIN;
    one[u] = LLONG_MIN; two[u] = LLONG_MIN;
    for (auto v: adj[u]) {
        // avoid going back to the parent node
        adj[v].erase(std::remove(adj[v].begin(), adj[v].end(), u), adj[v].end());
        dfs(v); // child's sum[v], one[v], two[v] are all ready to use after this.
        sum[u] += sum[v];
        one[u] = max(one[u], one[v]);
        two[u] = max(two[u], two[v]);
        // How to find the largest and the second largest element in an array in one pass?
        // https://www.techcrashcourse.com/2016/08/program-to-find-largest-and-second-largest-element-array.html
        if (one[v] >= one_firstmax) {
            one_secondmax = one_firstmax;
            one_firstmax = one[v];
        } else if (one[v] >= one_secondmax)
            one_secondmax = one[v];
    }
    sum[u] += a[u];
    one[u] = max(one[u], sum[u]);
    if (adj[u].size() >= 2) // should be equivalent to "one_secondmax > LLONG_MIN;"
        two[u] = max(two[u], one_firstmax + one_secondmax);
}

int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); // IO 優化
    cin >> n;
    for (int i=1; i<=n; i++) cin >> a[i];
    for (int i=1; i<n; i++) {
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dfs(1);
    if (two[1] == LLONG_MIN) cout << "Impossible";
    else cout << two[1];
    return 0;
}

結果： Example image