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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
|
#include <iostream>
#include <vector>
#include <sys/resource.h>
struct Vertex {
int weight;
std::vector <int> children;
};
typedef std::vector<Vertex> Graph;
typedef std::vector<int> Sum;
Graph ReadTree() {
int vertices_count;
std::cin >> vertices_count;
Graph tree(vertices_count);
for (int i = 0; i < vertices_count; ++i)
std::cin >> tree[i].weight;
for (int i = 1; i < vertices_count; ++i) {
int from, to, weight;
std::cin >> from >> to;
tree[from - 1].children.push_back(to - 1);
tree[to - 1].children.push_back(from - 1);
}
return tree;
}
int dfs(const Graph &tree, int vertex, int parent, std::vector<int> &w) {
if ( w[vertex] == -1) {
if (tree[vertex].children.size() == 0) {
w[vertex] = tree[vertex].weight;
} else {
int m0 = tree[vertex].weight;
for (int child : tree[vertex].children) {
if (child != parent) {
for (int gchild : tree[child].children) {
if (gchild != child && gchild != vertex) {
m0 += dfs(tree, gchild, child, w);
}
}
}
}
int m1 = 0;
for (int child : tree[vertex].children) {
if (child != parent) {
m1 += dfs(tree, child, vertex, w);
}
}
w[vertex] = std::max(m0, m1);
}
}
return w[vertex];
}
int MaxWeightIndependentTreeSubset(const Graph &tree) {
size_t size = tree.size();
if (size == 0)
return 0;
std::vector<int> w(size, -1);
return dfs(tree, 0, -1, w);
}
int main() {
// This code is here to increase the stack size to avoid stack overflow
// in depth-first search.
const rlim_t kStackSize = 64L * 1024L * 1024L; // min stack size = 64 Mb
struct rlimit rl;
int result;
result = getrlimit(RLIMIT_STACK, &rl);
if (result == 0)
{
if (rl.rlim_cur < kStackSize)
{
rl.rlim_cur = kStackSize;
result = setrlimit(RLIMIT_STACK, &rl);
if (result != 0)
{
fprintf(stderr, "setrlimit returned result = %d\n", result);
}
}
}
// Here begins the solution
Graph tree = ReadTree();
int weight = MaxWeightIndependentTreeSubset(tree);
std::cout << weight << std::endl;
return 0;
}
|