Update Count Nodes Equal to Average of Subtree.java

varunu28 · web-flow · commit 02b74a084498 · 2023-11-02T11:48:18.000+05:30
diff --git a/Medium/Count Nodes Equal to Average of Subtree.java b/Medium/Count Nodes Equal to Average of Subtree.java
@@ -14,24 +14,23 @@
  * }
  */
 class Solution {
-  public int averageOfSubtree(TreeNode root) {
-    int[] result = {0};
-    helper(root, result);
-    return result[0];
-  }
-  
-  private int[] helper(TreeNode root, int[] result) {
-    if (root == null) {
-      return new int[]{0, 0};
+    public int averageOfSubtree(TreeNode root) {
+        int[] result = {0};
+        recurse(root, result);
+        return result[0];
     }
-    int[] left = helper(root.left, result);
-    int[] right = helper(root.right, result);
-    int sum = left[0] + right[0] + root.val;
-    int numOfNodes = left[1] + right[1] + 1;
-    int average = sum / numOfNodes;
-    if (root.val == average) {
-      result[0]++;
+
+    private int[] recurse(TreeNode root, int[] result) {
+        if (root == null) {
+            return new int[]{0, 0};
+        }
+        int[] left = recurse(root.left, result);
+        int[] right = recurse(root.right, result);
+        int sum = left[0] + right[0] + root.val;
+        int count = left[1] + right[1] + 1;
+        if (count != 0 && sum / count == root.val) {
+            result[0]++;
+        }
+        return new int[]{sum, count};
     }
-    return new int[]{sum, numOfNodes};
-  }
 }
