leetcode

commit 5fe98805692cf3739cb4538098b178ba6c218223
parent d490d13063634c4aa9a2dfb12551ec0e5ff2d248
Author: Dimitrije Dobrota <mail@dimitrijedobrota.com>
Date:   Thu,  2 Nov 2023 22:45:05 +0000

Daily Problem

Diffstat:
M
Problems/2265.cpp
 | 
43
+++++++++++++++++++------------------------
M
README.md
 | 
1
+

2 files changed, 20 insertions(+), 24 deletions(-)
diff --git a/Problems/2265.cpp b/Problems/2265.cpp
@@ -1,33 +1,28 @@
 class Solution {
-  public:
-    int averageOfSubtree(TreeNode *root) {
-        stack<TreeNode *> st;
-        unordered_map<TreeNode *, int> um;
-        int res = 0;
+    typedef tuple<int, int> tii;
 
+  public:
+    int averageOfSubtree(const TreeNode *root) const {
+        unordered_map<const TreeNode *, tii> um;
+        stack<const TreeNode *> st;
+        int count = 0;
         st.push(root);
         while (!st.empty()) {
-            TreeNode *root = st.top();
-            if (root == nullptr) {
-                st.pop(), root = st.top(), st.pop();
-                int sum = root->val, count = 1;
-                if (root->left) {
-                    sum += root->left->val;
-                    count += um[root->left];
-                }
-                if (root->right) {
-                    sum += root->right->val;
-                    count += um[root->right];
-                }
-                if (root->val == sum / count) res++;
-                um[root] = count;
-                root->val = sum;
+            if (st.top() != nullptr) {
+                const TreeNode *root = st.top();
+                st.push(nullptr);
+                if (root->left) st.push(root->left);
+                if (root->right) st.push(root->right);
                 continue;
             }
-            st.push(nullptr);
-            if (root->left) st.push(root->left);
-            if (root->right) st.push(root->right);
+            st.pop();
+            const TreeNode *root = st.top();
+            st.pop();
+
+            tii left = um[root->left], right = um[root->right];
+            um[root] = {root->val + get<0>(left) + get<0>(right), 1 + get<1>(left) + get<1>(right)};
+            if (get<0>(um[root]) / get<1>(um[root]) == root->val) count++;
         }
-        return res;
+        return count;
     }
 };
diff --git a/README.md b/README.md
@@ -800,6 +800,7 @@ for solving problems.
 |  2246  |    Hard    | [Longest Path With Different Adjacent Characters](Problems/2246.cpp)                               |
 |  2251  |    Hard    | [Number of Flowers in Full Bloom](Problems/2251.cpp)                                               |
 |  2265  |   Medium   | [Count Nodes Equal to Average of Subtree](Problems/2265.cpp)                                       |
+|  2265  |   Medium   | [Count Nodes Equal to Average of Subtree](Problems/2265.cpp)                                       |
 |  2272  |    Hard    | [Substring With Largest Variance](Problems/2272.cpp)                                               |
 |  2275  |   Medium   | [Largest Combination With Bitwise AND Greater Than Zero](Problems/2275.cpp)                        |
 |  2279  |   Medium   | [Maximum Bags With Full Capacity of Rocks](Problems/2279.cpp)                                      |