LeetCode 222. Count Complete Tree Nodes
はじめに
LeetCode 「222. Count Complete Tree Nodes」の問題をGoで解きました。
問題
問題文
Given the root of a complete binary tree, return the number of the nodes in the tree.
According to Wikipedia, every level, except possibly the last, is completely filled in a complete binary tree, 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.
Design an algorithm that runs in less than O(n) time complexity.
和訳
完全二分木のルートが与えられたとき、その木のノードの数を返す。
ウィキペディアによると、完全二分木では最後のレベルを除くすべてのレベルが完全に埋まり、最後のレベルのすべてのノードは可能な限り左にある。最後のレベルhに含まれるノードは1～2hである。
O(n)未満の時間複雑度で実行されるアルゴリズムを設計せよ。
例
1
/ \
2 3
/ \ /
4 5 6
Input: root = [1,2,3,4,5,6]
Output: 6
Input: root = []
Output: 0
Input: root = [1]
Output: 1
制約
- The number of nodes in the tree is in the range [0, 5 * 104].
- 0 <= Node.val <= 5 * 104
- The tree is guaranteed to be complete.
解答
完全二分木のノード数を求める問題です。
完全二分木という前提がなければ、DFSで以下のように求めることができます。
コード
func countNodes(root *TreeNode) int {
if root == nil {
return 0
}
return countNodes(root.Left) + countNodes(root.Right) + 1
}
この実装だと、全てのノードをたどることになり、時間的計算量はO(n)となります。
今回の問題は、時間的計算量をO(n)未満で求める必要があるため、完全二分木の特徴を活かし、効率的に求める必要があります。
完全二分木
完全二分木は最下層以外は全ての層が埋まっており、最下層は左詰となっています。
そのため、左右の高さが同じであれば、左部分木は完全に埋まっている完全二分木と言えます。
完全に埋まった高さ h の完全二分木のノード数は 2^h - 1 です。
たとえば、高さが3なら 2^3 - 1 = 7 ノードとなります。
また、左右の高さが違う場合、右部分は1つ低い状態です。
その場合、右部分木は完全に埋まっている完全二分木となります。
コード
func countNodes(root *TreeNode) int {
if root == nil {
return 0
}
leftDepth := getDepth(root.Left)
rightDepth := getDepth(root.Right)
if leftDepth == rightDepth {
return (1 << leftDepth) + countNodes(root.Right)
} else {
return (1 << rightDepth) + countNodes(root.Left)
}
}
func getDepth(node *TreeNode) int {
depth := 0
for node != nil {
depth++
node = node.Left
}
return depth
}
(1 << leftDepth)の部分は、左部分木のノード数にルートノードを合わせた数を計算しています。
a << nは a × 2^n（2の累乗）という意味になります。
leftDepth == rightDepthの場合、
左部分木は高さleftDepthの完全二分木で、2^leftDepth - 1のノードがあります。
ルートを合わせると2^leftDepth個になるため、残りの右部分きに対して再帰的にcountNodes を呼び出します。
参考
二分木の問題は、以下の記事でも紹介していますので、参考にしてください。
Discussion