add a solution for 222

fishercoder1534 · fishercoder1534 · commit 9bc71e3727ca · 2021-10-25T11:48:25.000-07:00
diff --git a/src/main/java/com/fishercoder/solutions/_222.java b/src/main/java/com/fishercoder/solutions/_222.java
@@ -37,4 +37,69 @@ private int getLeftHeight(TreeNode root) {
         }
     }
 
+    public static class Solution2 {
+        /**
+         * credit: https://leetcode.com/problems/count-complete-tree-nodes/solution/
+         */
+        public int countNodes(TreeNode root) {
+            if (root == null) {
+                return 0;
+            }
+            int depth = getDepth(root);
+            if (depth == 0) {
+                return 1;
+            }
+            int left = 0;
+            int right = (int) Math.pow(2, depth) - 1;
+            //here the condition needs to be not bigger than right, instead of the typical strictly smaller than right.
+            while (left <= right) {
+                int mid = left + (right - left) / 2;
+                //this is to suppose the elements on the last level are numbered from 1 to Math.pow(2, d) - 1, we are using
+                //binary search here to find the right-most number
+                if (exists(root, mid, depth)) {
+                    left = mid + 1;
+                } else {
+                    right = mid - 1;
+                }
+            }
+            //number of all nodes equals all nodes in the previous level + all the nodes on the last level indicated by variable "left"
+            return (int) Math.pow(2, depth) - 1 + left;
+        }
+
+        private boolean exists(TreeNode root, int target, int depth) {
+            /**An example complete tree in this algorithm:
+             *         1
+             *        /  \
+             *       2    3
+             *      / \  /
+             *     1  2  3     (we use 1, 2, 3 at this level to indicate how this program runs instead of 4, 5, 6)
+             *
+             *     first, target is 1, we found 1 <= 1 (root), so we go to root.left, after going down to the last level (depth),
+             *     we found this target exists: node != null, we return true from this function
+             * */
+            int left = 0;
+            int right = (int) Math.pow(2, depth) - 1;
+            for (int i = 0; i < depth; i++) {
+                int mid = left + (right - left) / 2;
+                if (target <= mid) {
+                    root = root.left;
+                    right = mid;
+                } else {
+                    root = root.right;
+                    left = mid + 1;
+                }
+            }
+            return root != null;
+        }
+
+        private int getDepth(TreeNode root) {
+            int depth = 0;
+            while (root.left != null) {
+                root = root.left;
+                depth++;
+            }
+            return depth;
+        }
+    }
+
 }
diff --git a/src/test/java/com/fishercoder/_222Test.java b/src/test/java/com/fishercoder/_222Test.java
@@ -0,0 +1,45 @@
+package com.fishercoder;
+
+import com.fishercoder.common.classes.TreeNode;
+import com.fishercoder.common.utils.CommonUtils;
+import com.fishercoder.common.utils.TreeUtils;
+import com.fishercoder.solutions._2007;
+import com.fishercoder.solutions._222;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+
+public class _222Test {
+    private static _222.Solution1 solution1;
+    private static _222.Solution2 solution2;
+    private static int expected;
+    private static TreeNode root;
+
+    @BeforeClass
+    public static void setup() {
+        solution1 = new _222.Solution1();
+        solution2 = new _222.Solution2();
+    }
+
+    @Test
+    public void test1() {
+        root = TreeUtils.constructBinaryTree(Arrays.asList(1, 2, 3));
+        TreeUtils.printBinaryTree(root);
+        expected = 3;
+        assertEquals(expected, solution1.countNodes(root));
+        assertEquals(expected, solution2.countNodes(root));
+    }
+
+    @Test
+    public void test2() {
+        root = TreeUtils.constructBinaryTree(Arrays.asList(1, 2, 3, 4, 5, 6));
+        TreeUtils.printBinaryTree(root);
+        expected = 6;
+        assertEquals(expected, solution1.countNodes(root));
+        assertEquals(expected, solution2.countNodes(root));
+    }
+
+}
