1 files changed, 26 insertions, 15 deletions
diff --git a/02-data_structures/01-basic_data_structures/02-tree_height/tree-height.cpp b/02-data_structures/01-basic_data_structures/02-tree_height/tree-height.cpp
index d050e88..c6a55cd 100644
--- a/02-data_structures/01-basic_data_structures/02-tree_height/tree-height.cpp
+++ b/02-data_structures/01-basic_data_structures/02-tree_height/tree-height.cpp
@@ -3,28 +3,39 @@
 #include <algorithm>
 
 class TreeHeight {
-  int n;             
+  int n;
+  int root;
   std::vector<int> parent;
 
 public:
   void read() {
-    std::cin >> n;
-    parent.resize(n);
-    for (int i = 0; i < n; i++)
-      std::cin >> parent[i];
+      std::cin >> n;
+      parent.resize(n);
+      for (int i = 0; i < n; i++) {
+          std::cin >> parent[i];
+          if (parent[i] < 0)
+              root = i;
+      }
+      /* std::cout << root << std::endl; */
   }
 
-  int compute_height() {
-    // Replace this code with a faster implementation
-    int maxHeight = 0;
-    for (int vertex = 0; vertex < n; vertex++) {
-      int height = 0;
-      for (int i = vertex; i != -1; i = parent[i])
-        height++;
-      maxHeight = std::max(maxHeight, height);
-    }
-    return maxHeight;
+  int compute_height()
+  {
+      std::vector<int> heights(n);
+      for (int vertex = 0; vertex < n; vertex++) {
+          if (heights[vertex] != 0)
+              continue;
+          int height = 0;
+          for (int i = vertex; i != -1; i = parent[i]) {
+              height++;
+              if (height > heights[i])
+                  heights[i] = height;
+              else break;
+          }
+      }
+      return heights[root];
   }
+
 };
 
 int main() {