Count Complete Tree Nodes

1package pkg5leetcode.top100;2 3/*4 * Count Complete Tree Nodes | LC 2225 * APPROACH: Compare left/right heights; recurse one side if equal.6 * COMPLEXITY: Time O(log^2 n), Space O(log n)7 */8public class top100_LC222CountCompleteTreeNodes {9    /** Same shape as pkg5leetcode/common/TreeNode.java (nested for single-file runs). */10 11    static class TreeNode {12        int val;13        TreeNode left, right;14        TreeNode(int val) { this.val = val; }15    }16 17    static int countNodes(TreeNode root) {18        if (root == null) return 0;19        int left = depthLeft(root), right = depthRight(root);20        if (left == right) return (1 << left) - 1;21        return 1 + countNodes(root.left) + countNodes(root.right);22    }23 24    static int depthLeft(TreeNode n) {25        int d = 0;26        while (n != null) { d++; n = n.left; }27        return d;28    }29 30    static int depthRight(TreeNode n) {31        int d = 0;32        while (n != null) { d++; n = n.right; }33        return d;34    }35 36    public static void main(String[] args) {37        TreeNode root = new TreeNode(1);38        root.left = new TreeNode(2);39        root.right = new TreeNode(3);40        root.left.left = new TreeNode(4);41        root.left.right = new TreeNode(5);42        root.right.left = new TreeNode(6);43        check(countNodes(root) == 6, "case1");44        check(countNodes(null) == 0, "case2");45        System.out.println("all tests passed");46    }47 48    static void check(boolean cond, String name) {49        if (!cond) throw new AssertionError("FAILED: " + name);50        System.out.println("  PASS " + name);51    }52}