Обход бинарного дерева по слоям/уровням

Ссылка на задачу на leetcode: https://leetcode.com/problems/binary-tree-level-order-traversal/description/

Условие. Дано двоичное дерево. Нужно обойти его по слоям/уровням. Каждый уровень обходить слева направо. Результат вернуть в виде списка списков. Каждый вложенный список это элементы одного уровня.

Например:

Для такого дерева

Ответ: [[1], [2, 3], [3, 4, 5], [6]]

Решение. Для решения будем использовать поиск в ширину. Я описал алгоритм поиска в ширину тут: BFS.

Небольшая справка по деревьям.

Граф — это множество вершин и связей между ними.

Дерево — это связный граф, не содержащий циклов. Это древовидная структура, в которой есть корень, которой содержит ноль входящих ребер, а остальные вершины имеют ровно одно входящее ребро.

Двоичное дерево — дерево, в котором каждая вершина имеет не более двух дочерних вершин.

Класс вершины двоичного дерева можно описать на Java следующим образом:

class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;

  public TreeNode(int val, TreeNode left, TreeNode right) {
    this.val= val;
    this.left = left;
    this.right = right;
  }
}

left и right — это левый и правый ребенок вершины. Для простоты, на собеседованиях можно не объявлять поля private и не использовать геттеры для получения доступа к полям. Это позволит обращаться к полям как node.left, node.right.

Вернемся к задаче.

Как я уже сказал, мы будем использовать BFS. Для случая дерева его можно упростить. Нам не нужно использовать visited. Т.к. мы будем идти от родителей к детям. Вечного цикла у нас не возникнет.

Чтобы понять когда у нас начался и закончился уровень, который мы обходим, нам нужно хранить в очереди номер уровня или элементы каждого уровня складывать в отдельную очередь.

Рассмотрим второй вариант, с отдельной очередью для каждого уровня. Поместим корень дерева в очередь. Создадим еще одну очередь для следующего уровня. Будем итерироваться по первой очереди, пока она не закончится, но помещать детей не в ту же самую очередь, а во вторую. Как только первая очередь закончится — заменим ее второй и повторим весь процесс.

Решение на Java:

class Solution {
    public List<List<Integer>> levelOrder(TreeNode root) {
        List<List<Integer>> result = new ArrayList<>();
        Queue<TreeNode> queue = new LinkedList<>();
        if (root != null)
            queue.add(root);
        while (!queue.isEmpty()) {
            //Вторая очередь для следующего уровня
            Queue<TreeNode> newQueue = new LinkedList<>();
            //Тут будем хранить элементы одного уровня
            List<Integer> values = new ArrayList<>();
            //Цикл по первой очереди, но детей помещаем во вторую очередь
            while (!queue.isEmpty()) {
                TreeNode node = queue.poll();
                values.add(node.val);
                if (node.left != null)
                    newQueue.add(node.left);
                if (node.right != null)
                    newQueue.add(node.right);
            }
            result.add(values);
            //Подменяем первую очередь второй
            queue = newQueue;
        }
        return result;
    }
}