Added Count Complete Tree Nodes logNxlogN

Ace-Krypton · Sep 20, 2023 · 2718071 · 2718071
1 parent 5be8731
commit 2718071
Showing 1 changed file with 31 additions and 5 deletions.
diff --git a/Easy/CountCompleteTreeNodes/include/solution.hpp b/Easy/CountCompleteTreeNodes/include/solution.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cmath>
 #include <vector>
 #include <iostream>
 #include <algorithm>
@@ -16,16 +17,41 @@ struct TreeNode {
 class Solution {
 public:
     static auto count_nodes(TreeNode *root) -> int32_t {
+        if (!root) return 0;
+
+        std::function<int32_t(TreeNode*)> lh =
+                [&](TreeNode *root) -> int32_t {
+                    if (!root) return 0;
+                    return lh(root->left) + 1;
+                };
+
+        std::function<int32_t(TreeNode*)> rh =
+                [&](TreeNode *root) -> int32_t {
+                    if (!root) return 0;
+                    return rh(root->right) + 1;
+                };
+
+        const int32_t left_height = lh(root);
+        const int32_t right_height = rh(root);
+
+        if (left_height == right_height) {
+            return std::pow(2, left_height) - 1;
+        }
+
+        return count_nodes(root->left) + count_nodes(root->right) + 1;
+    }
+
+    [[maybe_unused]] static auto count_nodes_tle(TreeNode *root) -> int32_t {
         if (root == nullptr) return 0;
 
         int32_t count = 0;
         std::function<void(TreeNode*)> preorder =
                 [&](TreeNode *root) -> void {
-            if (root == nullptr) return;
-            ++count;
-            preorder(root->left);
-            preorder(root->right);
-        };
+                    if (root == nullptr) return;
+                    ++count;
+                    preorder(root->left);
+                    preorder(root->right);
+                };
 
         preorder(root);
         return count;