Binary Tree を BFS で探索する時、階層ごとに処理する。

leetCode の「Populating Next Right Pointers in Each Node」を解いた。
https://leetcode.com/problems/populating-next-right-pointers-in-each-node/

この問題を解く際、Binary Tree を BFS (幅優先探索) で探索しながら、階層ごとに処理するということを行った。この手法は今後も活用できそうだったので、ここメモしておく。

import java.util.ArrayDeque;
import java.util.Deque;

public class Solution {
    public Node connect(Node root) {
        if (root == null) {
            return null;
        }
        Deque<Node> queue = new ArrayDeque<>();
        queue.offer(root);

        while (!queue.isEmpty()) {
            // ★ 階層ごとに処理するため、ここで、キュー内の Node 数 (＝その階層の Node 数) を取得する。
            int count = queue.size();
            Node prevNode = null;
            while (count > 0) { // ★ その階層の Node 数の分だけ処理するためのループ。
                Node curNode = queue.poll();

                if (prevNode != null) {
                    prevNode.next = curNode;
                }
                prevNode = curNode;

                if (curNode.left != null) {
                    queue.offer(curNode.left);
                }
                if (curNode.right != null) {
                    queue.offer(curNode.right);
                }

                count--;
            }
        }
        return root;
    }
}

ポイントは、上記のソースのコメントの ★ を付けたところ。