Задача с собеседования в Google: Удалить вершины в дереве и вернуть оставшийся лес/forest

Задача.

Дано бинарное дерево. Все значения в вершинах уникальны в рамках всего дерева. Также дан массив значений, которые нужно удалить. Нужно удалить вершины с данными значениями. После удаления этих вершин у нас останутся разъединенные деревья (лес или forest). Нужно вернуть корни оставшихся после удаления деревьев в любом порядке.
Например,
для дерева:

После удаления вершин 3 и 5:

У нас останется лес из трех деревьев:

После произведения операций удаления, нужно вернуть корни этих деревьев: [1, 6, 7]

Ссылка на leetcode: https://leetcode.com/problems/delete-nodes-and-return-forest/

Решение.

Как и в 99% всех задач на бинарные деревья нам нужно, просто, применить алгоритм обхода бинарного дерева. (Алгоритмы обхода двоичного дерева)

При обходе дерева нам нужно сравнивать значение в вершине со значениями, которые нужно удалить.
Если в вершине значение, которое нам нужно удалить, то удаляем эту вершину. Но что значит удалить вершину? Нам нужно проапдейтить ссылку на эту вершину на null с родительской вершины.
Например, когда при обходе дерева мы дошли до вершины 3, она в списке на удаление. Чтобы ее удалить, нам нужно для ее родительской вершины удалить на нее ссылку:

parent_node.right = null

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

Стандартный алгоритм обхода дерева без дополнительной логики:

traverse(TreeNode root) {
    if (root == null) {
        return;
    }

    traverse(root.left);
    traverse(root.right);
}

Теперь давайте дополним этот код, чтобы работало удаление. Например, это можно сделать так: можно в качестве результата работы этой функции возвращать текущую вершину, если ее удалять не надо и возвращать null, если нужно удалить это вершину. Результат работы функции для левого ребенка будем присваивать root.left. Соотвественно, если она вернет null, то мы таким образом удалим ссылку с родительской вершины на удаленную вершину. Если она вернет текущую вершину, то такое присвоение ничего не изменит. Аналогично для правой вершины:

root.left = traverse(root.left);
root.right = traverse(root.right);

Тогда код по обходу дерева и удалению вершин будет таким:

TreeNode traverse(TreeNode root, Set<Integer> valuesToDelete) {
    if (root == null) {
        return null;
    }
    //traverse вернет null, если дочерняя вершина удалена
    //это позволит удалить ссылку с родительской вершины 
    //на дочернюю
    root.left = traverse(root.left, valuesToDelete);
    root.right = traverse(root.right, valuesToDelete);

    //Если значение в вершине в списке значений на удаление, 
    //то вернем null, чтобы удалить ссылку с родительской 
    //вершины на удаленную вершину
    if (valuesToDelete.contains(root.val)) {
        return null;
    }
    //Иначе вернем текущую вершину, 
    //и ссылка с родительской вершины не изменится
    return root;
}

Весь код, включая вызов функции обхода дерева будет выглядеть так:

public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
    //Результат. Пока оставим пустым, 
    //позже добавим логику заполнения
    List<TreeNode> result = new ArrayList<>();
    //Преобразуем массив в Set, для быстрого доступа за O(1)
    Set<Integer> valuesToDelete = new HashSet<>();
    for (int val : to_delete) {
        valuesToDelete.add(val);
    }
    //Вызов функции обхода дерева
    traverse(root, valuesToDelete, result);
    return result;
}

private TreeNode traverse(TreeNode root, Set<Integer> valuesToDelete, List<TreeNode> result) {
    if (root == null) {
        return null;
    }

    root.left = traverse(root.left, valuesToDelete, result);
    root.right = traverse(root.right, valuesToDelete, result);

    if (valuesToDelete.contains(root.val)) {
        return null;
    }
    return root;
}

Теперь добавим логику для заполнения result. Туда нужно помещать новые корни деревьев, после удаления.
Тут логика простая. Когда мы удаляем вершину, то ее дети становятся новыми корнями. Поэтому их надо добавить в качестве результата. Например, при удалении вершины 3, нам нужно добавить в результат вершины 6 и 7, т.к. они станут новыми корнями деревьев в нашем лесу:

Соответственно, добавим внутрь условия проверки, что текущая вершина должна быть удалена, код добавления ее детей в результат:

    if (valuesToDelete.contains(root.val)) {
        //Добавляем детей вершины, которую удаляем в результат
        if (root.left != null) {
            result.add(root.left);
        }
        if (root.right != null) {
            result.add(root.right);
        }
        return null;
    }

Но тут можно подумать, но, что если 6 и 7 тоже нужно удалить, а мы их в результат добавляем?
В силу того, что мы делаем post-order обход, мы вначале посещаем всех детей, а потом применяем логику на родительской вершине. Если 6 или 7 тоже нужно удалить, то мы обновим ссылки root.left и root.right раньше. И они уже будут равными null, поэтому в результат они добавлены не будут.

Также тут есть edge-case для изначального корня заданного дерева. Если его удалять не надо, то его надо также добавить в результат, иначе наша логика не покроект этот случай:

    if (!valuesToDelete.contains(root.val)) {
        result.add(root);
    }

Все вместе:

public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
    List<TreeNode> result = new ArrayList<>();
    //Преобразуем массив в Set, для быстрого доступа за O(1)
    Set<Integer> valuesToDelete = new HashSet<>();
    for (int val : to_delete) {
        valuesToDelete.add(val);
    }
    //Edge-case. Если корень изначального дерева не надо удалять,
    //то добавляем его в результат
    if (!valuesToDelete.contains(root.val)) {
        result.add(root);
    }
    //Вызываем функцию обхода дерева
    traverse(root, valuesToDelete, result);
    return result;
}

private TreeNode traverse(TreeNode root, Set<Integer> valuesToDelete, List<TreeNode> result) {
    if (root == null) {
        return null;
    }
    //traverse вернет null, если дочерняя вершина удалена
    //это позволит удалить ссылку с родительской вершины 
    //на дочернюю
    root.left = traverse(root.left, valuesToDelete, result);
    root.right = traverse(root.right, valuesToDelete, result);

    if (valuesToDelete.contains(root.val)) {
        //Добавляем детей вершины, которую удаляем в результат
        if (root.left != null) {
            result.add(root.left);
        }
        if (root.right != null) {
            result.add(root.right);
        }
        //Если значение в вершине в списке значений на удаление, 
        //то вернем null, чтобы удалить ссылку с родительской 
        //вершины на удаленную вершину
        return null;
    }
    return root;
}

Time Complexity = O(N), Space Complexity = O(N). Где N - число вершин в дереве. Мы обходим все вершины дерева. Space Complexity берется из того, что у нас рекурсивные вызовы и у нас расходуется память на call stack. В худшем случае дерево может быть в виде одной линии, глубины N. Поэтому погружение будет в худшем случае на N. Для сбалансированного дерева Space Complexity = O(log(N)), log(N) будет высотой сбалансированного дерева. Т.к. в условии ничего про сбалансированность дерева нет, то Space Complexity = O(N).