問題
二分木のroot を
- 各節は
val(値)、left(左の子)、 right(右の子) を 持ちます。 子が 無い ところは、 Python では None、JavaScript では nullです - 木が
空 ( rootが無い ) の ときは、 空の 並び []を返します - 節の
数は 0 以上 2000 以下、 値は どれも -1000 以上 1000 以下です
例では、null です。
入力と出力の例
例 1
根が 5、
入力
root = [5, 2, 8, 1, null, 6, 9]出力
[[5], [2, 8], [1, 6, 9]]例 2
右へ、
入力
root = [4, null, 7, 3]出力
[[4], [7], [3]]例 3
空の木です。
入力
root = []出力
[]面接での聞かれ方
面接では、
- 「二分木を、
段ごとに 値を まとめて 返してください」 - 書き終えた
あとに 「右から 読む形に できますか？」 「段ごとではなく、 ジグザグに 読めますか？」と 少し ずつ形を 変えて 聞かれます - 「再帰でも書けますか？」と、
別の 書き方を 聞かれる こともあります
木の
まず
段 1
まずは、
段 2
段の
段 3
最後に、
段 4
待ち行列のcollections.deque を
ここから解答です。先に自分で書いてから読むと、考え方が残ります。
解答
Python
Python
from collections import deque
def tree_levels(root):
if root is None:
return []
levels = []
queue = deque([root])
add = queue.append
while queue:
row = []
size = len(queue)
for _ in range(size):
nd = queue.popleft()
row.append(nd.val)
if nd.left:
add(nd.left)
if nd.right:
add(nd.right)
levels.append(row)
return levelsJavaScript
JavaScript
function treeLevels(root) {
if (root === null) return [];
const levels = [];
const queue = [root];
let head = 0;
while (head < queue.length) {
const end = queue.length;
const row = [];
while (head < end) {
const node = queue[head];
head++;
row.push(node.val);
if (node.left) {
queue.push(node.left);
}
if (node.right) {
queue.push(node.right);
}
}
levels.push(row);
}
return levels;
}JavaScript の end は、end までで
計算量
- 時間は O(n) です。
どの 節も 待ち行列に 1 回入り、 1 回出るだけです - 使う
記憶は O(n) です。 いちばん節の 多い 段が、 待ち行列に 一度に 並びます。 返す 並びも 全部の 値を持ちます - JavaScript の
書き方は、 取り出した 節も 配列に 残したままにします。 その分の 記憶は 増えますが、 O(n) のままです
よく
- 配列の
先頭を 消して 取り出す。 JavaScript の shift()や Python のlist.pop(0)は、残りを 全部 1 つずつ前へずらします。 節が 多いと、 全体が O(n²) になります - 読みながら長さを見直す。
段の 長さを、 取り出すたびに len(queue)で見直すと、 足したばかりの 次の 段まで 同じ 段に 入ってしまいます。 読み始める 前に 1 回だけ数えます - 空の
木を 見落とす。 rootが無いのに 待ち行列へ 入れると、 取り出した 節の valを読もうとして 止まります。 最初に 空の 並びを 返しておきます
面接での
書き始める
「段の
順に 読みたいので、 深さ優先ではなく、 待ち行列を 使った 幅優先で 書きます。 段の 切れ目は、 その段を 読み始める 前の 待ち行列の 長さで 分けます。 どの 節も 1 回ずつしか 出し 入れしないので、 時間は O(n) です。」
書き終えたら、
「例 1 で
たどります。 最初は 5 だけが 入っていて、 長さ 1 なので 1 つ 取り出して [5]です。その とき 2 と 8 を 足します。 次は 長さ 2 なので 2 つ 取り出して [2, 8]になり、 1、 6、 9 が 足されます。 最後に 長さ 3 で [1, 6, 9]となり、 待ち行列が 空に なって 終わります。」