updates count nodes complete tree

Author: anishLearnsToCode

Diff for: src/CountCompleteTreeNodes.java

@@ -3,17 +3,46 @@
 // S: O(log(n))
 
 public class CountCompleteTreeNodes {
-    private static int depth(TreeNode root) {
-        if (root == null) return 0;
-        return 1 + depth(root.left);
+    public int countNodes(TreeNode root) {
+        if (root == null) {
+            return 0;
+        }
+
+        final int depth = depthTree(root) - 1;
+        if (depth == 0) {
+            return 1;
+        }
+
+        int left = 0, right = (int) Math.pow(2, depth) - 1, middle;
+        while (left <= right) {
+            middle = left + (right - left) / 2;
+            if (exists(root, depth, middle)) left = middle + 1;
+            else right = middle - 1;
+        }
+
+        return (int) Math.pow(2, depth) - 1 + left;
     }
 
-    public static int countNodes(TreeNode root) {
-        if (root == null) return 0;
-        int leftDepth = depth(root.left);
-        int rightDepth = depth(root.right);
-        return 1 + (leftDepth == rightDepth
-                ? ((1 << leftDepth) - 1) + countNodes(root.right)
-                : countNodes(root.left) + countNodes(root.right));
+    private static boolean exists(TreeNode root, int depth, int index) {
+        TreeNode current = root;
+        int value = -1;
+        for (int i = 0 ; i < depth ; i++) {
+            final int middle = (int) Math.pow(2, depth - 1 - i);
+            if (index > value + middle) {
+                current = current.right;
+                value += middle;
+            } else {
+                current = current.left;
+            }
+        }
+        return current != null;
+    }
+
+    private static int depthTree(TreeNode root) {
+        if (root == null) {
+            return 0;
+        }
+
+        return 1 + Math.max(depthTree(root.left), depthTree(root.right));
     }
 }

Diff for: src/HelloWorld.java

@@ -1,25 +1,44 @@
 public class HelloWorld {
-    private static int result = Integer.MIN_VALUE;
-
-    public static int maxPathSum(TreeNode root) {
-        result = Integer.MIN_VALUE;
+    public int countNodes(TreeNode root) {
         if (root == null) {
             return 0;
         }
-        maxSum(root);
-        return result;
+
+        final int depth = depthTree(root) - 1;
+        if (depth == 0) {
+            return 1;
+        }
+
+        int left = 0, right = (int) Math.pow(2, depth) - 1, middle;
+        while (left <= right) {
+            middle = left + (right - left) / 2;
+            if (exists(root, depth, middle)) left = middle + 1;
+            else right = middle - 1;
+        }
+
+        return (int) Math.pow(2, depth) - 1 + left;
     }
 
-    private static int maxSum(TreeNode root) {
+    private static boolean exists(TreeNode root, int depth, int index) {
+        TreeNode current = root;
+        int value = -1;
+        for (int i = 0 ; i < depth ; i++) {
+            final int middle = (int) Math.pow(2, depth - 1 - i);
+            if (index > value + middle) {
+                current = current.right;
+                value += middle;
+            } else {
+                current = current.left;
+            }
+        }
+        return current != null;
+    }
+
+    private static int depthTree(TreeNode root) {
         if (root == null) {
             return 0;
         }
 
-        final int maxGainLeftSubtree = Math.max(maxSum(root.left), 0);
-        final int maxGainRightSubtree = Math.max(maxSum(root.right), 0);
-        final int maxPathSum = root.val + maxGainLeftSubtree + maxGainRightSubtree;
-        result = Math.max(result, maxPathSum);
-
-        return Math.max(root.val + maxGainLeftSubtree, root.val + maxGainRightSubtree);
+        return 1 + Math.max(depthTree(root.left), depthTree(root.right));
     }
 }