algorithm-ps/2213.cpp at master · sky7th/algorithm-ps · GitHub

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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
#include <bits/stdc++.h>
using namespace std;

int weight[10000], cache[10000][2], N;
vector<int> adj[10000], result;
bool visited[10000], selected[10000];

int func(int root, int select) {
    int &ret = cache[root][select];
    if (ret != -1) return ret;
    visited[root] = true;
    int include = weight[root];
    int exclude = 0;
    for (int next : adj[root]) {
        if (!visited[next]) {
            include += func(next, false);
            exclude += func(next, true);
        }
    }
    visited[root] = false;  // ㅠㅡㅜ
    if (select && include > exclude) {
        selected[root] = true;
        return ret = include;
    }
    return ret = exclude;
}

void tracker(int root, bool select) {
    visited[root] = true;
    if (select && selected[root]) {
        result.push_back(root);
        for (int next : adj[root])
            if (!visited[next]) tracker(next, false);
    } else {
        for (int next : adj[root])
            if (!visited[next]) tracker(next, true);
    }
    visited[root] = false;
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cin >> N;
    memset(cache, -1, sizeof(cache));
    for (int i = 0; i < N; i++)
        cin >> weight[i];
    for (int i = 0; i < N - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u - 1].push_back(v - 1);
        adj[v - 1].push_back(u - 1);
    }
    cout << func(0, true) << '\n';
    tracker(0, true);
    sort(result.begin(), result.end());
    for (int node : result)
        cout << node + 1 << ' ';
}