diff --git a/Easy/CountCompleteTreeNodes/include/solution.hpp b/Easy/CountCompleteTreeNodes/include/solution.hpp index ebe5a10..59641c2 100644 --- a/Easy/CountCompleteTreeNodes/include/solution.hpp +++ b/Easy/CountCompleteTreeNodes/include/solution.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -16,16 +17,41 @@ struct TreeNode { class Solution { public: static auto count_nodes(TreeNode *root) -> int32_t { + if (!root) return 0; + + std::function lh = + [&](TreeNode *root) -> int32_t { + if (!root) return 0; + return lh(root->left) + 1; + }; + + std::function 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 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;