1110. Delete Nodes And Return Forest

leetcode

Links

Submission

Problem:

Given the `root` of a binary tree, each node in the tree has a distinct value. After deleting certain nodes, we need to return the forest (a collection of trees). Each node in the tree that needs to be deleted will be included in an array `to_delete`.

We need to:

Delete the nodes listed in `to_delete`
Return the remaining trees in the forest.

Solution Description:

To implement the problem we need to perform the following steps:

Perform a Depth-First Search (DFS) on the tree.
If a node needs to be deleted, recursively call the function on its children and add those children to the forest (if they are not null).
For each node, check if it's in the `to_delete` list. If it is, mark it for deletion and set its parent's reference to it to `null`.
Keep a record of all roots of trees that are not marked for deletion and return them at the end.

Time Complexity:

Each node is visited once, resulting in O(N) time complexity, where N is the number of nodes in the tree.

Space Complexity:

In the worst case scenario, the recursion stack can go as deep as the height of the tree (O(H)), but since we store the nodes in to_delete in a set and the result in a list, the space complexity is O(N).

Example:

Consider this binary tree:

If the `to_delete` list is [3, 5], the forest after deletion would be:

1     6
/ \
2  4  7

Setup:

We will define the `mainSolution` function to perform the DFS traversal for deleting the nodes and then recording the forest. We will also include a test setup to validate the implementation.

Test Execution:

The testing framework will display input parameters, the actual result, the expected result, and whether the test passed or failed for each test case.

Implementation:

// Helper functions for logging and debugging (no-op if NestedInteger is defined)

/**
 * Main solution function that performs tree manipulation
 * @param {TreeNode} root - The root of the binary tree
 * @param {number[]} to_delete - List of node values to be deleted
 * @return {TreeNode[]} - Forest of remaining trees
 */
function mainSolution(root, to_delete) {
    const toDeleteSet = new Set(to_delete);
    const forest = [];

    function dfs(node, isRoot) {
        if (!node) return null;

        const shouldDelete = toDeleteSet.has(node.val);
        if (isRoot && !shouldDelete) {
            forest.push(node);
        }

        node.left = dfs(node.left, shouldDelete);
        node.right = dfs(node.right, shouldDelete);

        return shouldDelete ? null : node;
    }

    dfs(root, true);
    return forest;
}

// Definition for a binary tree node.
class TreeNode {
    constructor(val = 0, left = null, right = null) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

// Test cases
const testCases = [
    {
        root: new TreeNode(1,
            new TreeNode(2, new TreeNode(4), new TreeNode(5)),
            new TreeNode(3, new TreeNode(6), new TreeNode(7))
        ),
        to_delete: [3, 5],
        expected: [ [1, 2, 4], [6], [7] ]
    },
    {
        root: new TreeNode(1, new TreeNode(2), new TreeNode(3)),
        to_delete: [1],
        expected: [[2], [3]],
    },
    {
        root: new TreeNode(1, new TreeNode(2, new TreeNode(4), new TreeNode(3)), new TreeNode(5)),
        to_delete: [2, 3],
        expected: [[1, 5], [4]],
    },
    {
        root: new TreeNode(1, new TreeNode(2, null, new TreeNode(3)), null),
        to_delete: [2],
        expected: [[1], [3]],
    },
    {
        root: new TreeNode(1, null, new TreeNode(2, new TreeNode(3), new TreeNode(4))),
        to_delete: [2],
        expected: [[1], [3], [4]],
    }
];

// Execute test cases
testCases.forEach((test, index) => {
    const result = mainSolution(test.root, test.to_delete);
    const resultFormatted = result.map(tree => flatten(tree));
    console.log(`Test Case ${index + 1}: ${arraysEqual(resultFormatted, test.expected) ? 'Passed' : 'Failed'} (Expected: ${JSON.stringify(test.expected)}, Got: ${JSON.stringify(resultFormatted)})`);
});

/**
 * Function to flatten a tree to an array
 * @param {TreeNode} node - The root node of the tree to flatten
 * @return {number[]} - Flattened tree as array
 */
function flatten(node) {
    if (!node) return [];
    return [node.val, ...flatten(node.left), ...flatten(node.right)];
}

/**
 * Function to check if two nested arrays are equal
 * @param {Array} a - First array
 * @param {Array} b - Second array
 * @return {boolean} - True if arrays are equal, false otherwise
 */
function arraysEqual(a, b) {
    if (a.length !== b.length) return false;
    for (let i = 0; i < a.length; i++) {
        if (Array.isArray(a[i]) && Array.isArray(b[i])) {
            if (!arraysEqual(a[i], b[i])) return false;
        } else if (a[i] !== b[i]) {
            return false;
        }
    }
    return true;
}