222. Count Complete Tree Nodes

228 查看

题目:
Given a complete binary tree, count the number of nodes.

Definition of a complete binary tree from Wikipedia:
In a complete binary tree every level, except possibly the last, is completely filled, and all nodes in the last level are as far left as possible. It can have between 1 and 2h nodes inclusive at the last level h.

解答:

public class Solution {
    //学了enum的方法哈哈
    public enum Direction {
        LEFT, RIGHT
    }
    
    public int getDepth(TreeNode root, Direction dir) {
        int depth = 0;
        //这里是从root开始算起,所以求出来的结果是实际depth + 1
        while (root != null) {
            depth++;
            if (dir == Direction.LEFT) {
                root = root.left;
            } else {
                root = root.right;
            }
        }
        return depth;
    }
    
    public int countNodes(TreeNode root) {
        if (root == null) return 0;
        int leftDepth = getDepth(root, Direction.LEFT);
        int rightDepth = getDepth(root, Direction.RIGHT);
        
        if (leftDepth == rightDepth) {
            //1 << leftDepth是2^(leftDepth)
            return (1 << leftDepth) - 1;
        } else {
            //这一步很巧妙,把根结点加上,然后分治左结点和右结点,如果是叶子结点,在countNodes中就返回1
            return 1 + countNodes(root.left) + countNodes(root.right);
        }
    }
}