222. 完全二叉树的节点个数

简单二叉树递归 + 比较左右子树高度力扣原题 ↗

题目描述

输入一棵完全二叉树的根节点 root，返回树中节点的总数。

完全二叉树指：除最后一层外，每一层都排满；最后一层的节点从左往右连续排列，中间没有空缺。若把最后一层记作第 h 层（根是第 0 层），这一层有 1 到 2^h 个节点。

树以层序数组表示，例如 [1,2,3,4,5,6]；空树为 []。

示例 1

输入：root = [1,2,3,4,5,6]
输出：6

解释：树共有 6 个节点。

示例 2

输入：root = []
输出：0

解释：空树没有节点。

示例 3

输入：root = [1]
输出：1

解释：只有根节点。

数据范围

节点数在 [0, 5 * 10^4] 范围内。
0 <= Node.val <= 5 * 10^4
输入保证是一棵完全二叉树。

进阶：逐个遍历节点的做法是 O(n) 的。能否设计出更快的算法？

解题思路：递归 + 比较左右子树高度

分别沿左子树和右子树的最左链往下走，得到它们的高度 lh、rh。如果 lh == rh，说明最后一层已延伸到右子树，左子树是高 lh 的满二叉树，节点数 2^lh - 1，再加上根就是 2^lh，只需递归数右子树；否则右子树是高 rh 的满二叉树，计 2^rh 后递归数左子树。每层只递归一侧。

为什么正确

完全二叉树的任一子树仍是完全二叉树，并且左右子树中至少有一个是满二叉树。

复杂度分析

时间 O(log^2 n)：递归 O(log n) 层，每层量高度 O(log n)。空间 O(log n)。

易错点

量右子树高度时也要沿“最左”链走，因为完全二叉树最底层靠左，最左链长度才是子树高度。另外 1 << lh 已经包含了根节点本身，不要再加 1。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        if not root:
            return 0
        lh, node = 0, root.left
        while node:
            lh += 1
            node = node.left
        rh, node = 0, root.right
        while node:
            rh += 1
            node = node.left
        if lh == rh:
            return (1 << lh) + self.countNodes(root.right)
        return (1 << rh) + self.countNodes(root.left)

class Solution {
    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int lh = 0; TreeNode node = root.left;
        while (node != null) {
            lh++;
            node = node.left;
        }
        int rh = 0; node = root.right;
        while (node != null) {
            rh++;
            node = node.left;
        }
        if (lh == rh) {
            return (1 << lh) + countNodes(root.right);
        }
        return (1 << rh) + countNodes(root.left);
    }
}

class Solution {
public:
    int countNodes(TreeNode* root) {
        if (!root) {
            return 0;
        }
        int lh = 0; TreeNode* node = root->left;
        while (node) {
            lh++;
            node = node->left;
        }
        int rh = 0; node = root->right;
        while (node) {
            rh++;
            node = node->left;
        }
        if (lh == rh) {
            return (1 << lh) + countNodes(root->right);
        }
        return (1 << rh) + countNodes(root->left);
    }
};

func countNodes(root *TreeNode) int {
	if root == nil {
		return 0
	}
	lh, node := 0, root.Left
	for node != nil {
		lh++
		node = node.Left
	}
	rh := 0
	node = root.Right
	for node != nil {
		rh++
		node = node.Left
	}
	if lh == rh {
		return (1 << lh) + countNodes(root.Right)
	}
	return (1 << rh) + countNodes(root.Left)
}

var countNodes = function(root) {
    if (!root) {
        return 0;
    }
    let lh = 0, node = root.left;
    while (node) {
        lh++;
        node = node.left;
    }
    let rh = 0; node = root.right;
    while (node) {
        rh++;
        node = node.left;
    }
    if (lh === rh) {
        return (1 << lh) + countNodes(root.right);
    }
    return (1 << rh) + countNodes(root.left);
};

function countNodes(root: TreeNode | null): number {
    if (!root) {
        return 0;
    }
    let lh = 0, node: TreeNode | null = root.left;
    while (node) {
        lh++;
        node = node.left;
    }
    let rh = 0; node = root.right;
    while (node) {
        rh++;
        node = node.left;
    }
    if (lh === rh) {
        return (1 << lh) + countNodes(root.right);
    }
    return (1 << rh) + countNodes(root.left);
}

载入编辑器，自己跑一遍 →

其他解法

主解法每层比较左右子树高度，确定哪一边是满二叉树，然后递归另一边；二分解法把问题转化为“在最后一层找最右侧存在的结点”，用编号二分。两者都是 O(log² n)。

解法 2：二分查找 + 位运算时间 O(log² n)，空间 O(1)

先一直向左走，得到树的最大层号 h（根为第 0 层）。完全二叉树的结点数在 [2^h, 2^(h+1)-1] 之间。把结点按层序从 1 开始编号，编号 k 的二进制去掉最高位后，从高到低每一位就是从根走到它的路径：0 向左，1 向右。用这个办法可以在 O(h) 内判断编号 k 的结点是否存在。在范围内二分查找最大的存在编号，它就是结点总数。

不用递归，空间 O(1)，展示了“堆式编号的二进制位就是路径”这一重要技巧。边界和位运算细节较多，比主解法难写对。

from typing import Optional

class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        if root is None:
            return 0
        h, node = 0, root
        while node.left:
            h += 1
            node = node.left

        def exists(k):
            bit = 1 << (h - 1)
            node = root
            while node and bit:
                node = node.right if k & bit else node.left
                bit >>= 1
            return node is not None

        lo, hi = 1 << h, (1 << (h + 1)) - 1
        while lo < hi:
            mid = (lo + hi + 1) // 2
            if exists(mid):
                lo = mid
            else:
                hi = mid - 1
        return lo

class Solution {
    private boolean exists(TreeNode root, int h, int k) {
        int bit = h > 0 ? 1 << (h - 1) : 0;
        TreeNode node = root;
        while (node != null && bit > 0) {
            node = (k & bit) != 0 ? node.right : node.left;
            bit >>= 1;
        }
        return node != null;
    }

    public int countNodes(TreeNode root) {
        if (root == null) return 0;
        int h = 0;
        for (TreeNode node = root; node.left != null; node = node.left) h++;
        int lo = 1 << h, hi = (1 << (h + 1)) - 1;
        while (lo < hi) {
            int mid = lo + (hi - lo + 1) / 2;
            if (exists(root, h, mid)) lo = mid;
            else hi = mid - 1;
        }
        return lo;
    }
}

class Solution {
    bool exists(TreeNode* root, int h, int k) {
        int bit = h > 0 ? 1 << (h - 1) : 0;
        TreeNode* node = root;
        while (node && bit) {
            node = (k & bit) ? node->right : node->left;
            bit >>= 1;
        }
        return node != nullptr;
    }
public:
    int countNodes(TreeNode* root) {
        if (!root) return 0;
        int h = 0;
        for (TreeNode* node = root; node->left; node = node->left) h++;
        int lo = 1 << h, hi = (1 << (h + 1)) - 1;
        while (lo < hi) {
            int mid = lo + (hi - lo + 1) / 2;
            if (exists(root, h, mid)) lo = mid;
            else hi = mid - 1;
        }
        return lo;
    }
};

func countNodes(root *TreeNode) int {
	if root == nil {
		return 0
	}
	h := 0
	for node := root; node.Left != nil; node = node.Left {
		h++
	}
	exists := func(k int) bool {
		bit := 0
		if h > 0 {
			bit = 1 << (h - 1)
		}
		node := root
		for node != nil && bit > 0 {
			if k&bit != 0 {
				node = node.Right
			} else {
				node = node.Left
			}
			bit >>= 1
		}
		return node != nil
	}
	lo, hi := 1<<h, (1<<(h+1))-1
	for lo < hi {
		mid := lo + (hi-lo+1)/2
		if exists(mid) {
			lo = mid
		} else {
			hi = mid - 1
		}
	}
	return lo
}

var countNodes = function(root) {
    if (!root) return 0;
    let h = 0;
    for (let node = root; node.left; node = node.left) h++;
    const exists = (k) => {
        let bit = h > 0 ? 1 << (h - 1) : 0;
        let node = root;
        while (node && bit) {
            node = (k & bit) ? node.right : node.left;
            bit >>= 1;
        }
        return node !== null;
    };
    let lo = 1 << h, hi = (1 << (h + 1)) - 1;
    while (lo < hi) {
        const mid = lo + Math.floor((hi - lo + 1) / 2);
        if (exists(mid)) lo = mid;
        else hi = mid - 1;
    }
    return lo;
};

function countNodes(root: TreeNode | null): number {
    if (!root) return 0;
    let h = 0;
    for (let node: TreeNode = root; node.left; node = node.left) h++;
    const exists = (k: number): boolean => {
        let bit = h > 0 ? 1 << (h - 1) : 0;
        let node: TreeNode | null = root;
        while (node && bit) {
            node = (k & bit) ? node.right : node.left;
            bit >>= 1;
        }
        return node !== null;
    };
    let lo = 1 << h, hi = (1 << (h + 1)) - 1;
    while (lo < hi) {
        const mid = lo + Math.floor((hi - lo + 1) / 2);
        if (exists(mid)) lo = mid;
        else hi = mid - 1;
    }
    return lo;
};

二叉树专题的其他题目

查看二叉树全部题目 →