1 files changed, 28 insertions, 11 deletions
diff --git a/05-advanced_algorithms_and_complexity/04-np-completeness/02-plan_party/plan_party.cpp b/05-advanced_algorithms_and_complexity/04-np-completeness/02-plan_party/plan_party.cpp
index 70249aa..f612252 100644
--- a/05-advanced_algorithms_and_complexity/04-np-completeness/02-plan_party/plan_party.cpp
+++ b/05-advanced_algorithms_and_complexity/04-np-completeness/02-plan_party/plan_party.cpp
@@ -1,5 +1,6 @@
 #include <iostream>
 #include <vector>
+#include <sys/resource.h>
 
 struct Vertex {
     int weight;
@@ -27,23 +28,39 @@ Graph ReadTree() {
     return tree;
 }
 
-void dfs(const Graph &tree, int vertex, int parent) {
-    for (int child : tree[vertex].children)
-        if (child != parent)
-            dfs(tree, child, vertex);
-
-    // This is a template function for processing a tree using depth-first search.
-    // Write your code here.
-    // You may need to add more parameters to this function for child processing.
+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;
-    dfs(tree, 0, -1);
-    // You must decide what to return.
-    return 0;
+    std::vector<int> w(size, -1);
+    return dfs(tree, 0, -1, w);
 }
 
 int main() {