summaryrefslogtreecommitdiffstats
path: root/05-advanced_algorithms_and_complexity/04-np-completeness/02-plan_party/plan_party.cpp
blob: f6122523627f6ef797a90ec354dab5f9cd0f9d91 (plain)
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;
}