[leetcode] 1110. Delete Nodes And Return Forest

原題目

Given the root of a binary tree, each node in the tree has a distinct value.

After deleting all nodes with a value in to_delete, we are left with a forest (a disjoint union of trees).

Return the roots of the trees in the remaining forest. You may return the result in any order.

Example 1:

Input: root = [1,2,3,4,5,6,7], to_delete = [3,5]
Output: [[1,2,null,4],[6],[7]]
複製代碼

解答

這道題目須要使用 dfs（深度優先搜索）算法。一開始本身寫的 js 代碼是：

/** * Definition for a binary tree node. * function TreeNode(val) { * this.val = val; * this.left = this.right = null; * } */
/** * @param {TreeNode} root * @param {number[]} to_delete * @return {TreeNode[]} */
var delNodes = function(root, to_delete) {
    const res = new Map([[root, 1]])
    dfs(root, to_delete, res)
    return [...res.keys()];
};

const dfs = (node, to_delete, res, parent, type) => {
    const { left, right, val } = node
    const index = to_delete.indexOf(val)
    if (index > -1) {
        left && res.set(left, 1)
        right && res.set(right, 1)
        res.delete(node)
        if (type == 'left') {
            parent.left = null
        } else if (type == 'right') {
            parent.right = null
        }
    }
    left && dfs(left, to_delete, res, node, 'left')
    right && dfs(right, to_delete, res, node, 'right')
}
複製代碼

我使用了 es6 的 Map 數據結構，將須要保存在 forest 的節點做爲 Map 的 key 保存起來，而後當這個節點的 val 和 to_delete 的元素相等時，再將這個節點從 map 中 delete 掉。

但其實這個先 res.set 而後根據須要 res.delete 的操做顯得冗餘，同時須要作判斷並設置 prarent.left 和 parent.right，這樣也不太完美。應該有辦法能夠更加簡潔的實現。

參考了大神的代碼：

# Definition for a binary tree node.
# class TreeNode(object):
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None

class Solution(object):
    def delNodes(self, root, to_delete):
        to_delete_set = set(to_delete)
        res = []

        def helper(root, is_root):
            if not root: return None
            root_deleted = root.val in to_delete_set
            # 根據條件判斷是否加入，而非先加後刪
            if is_root and not root_deleted:
                res.append(root)
            # 若是這個節點被刪了，那麼它的子節點就是 forest 中的 root 了，合理！
            root.left = helper(root.left, root_deleted)
            root.right = helper(root.right, root_deleted)
            return None if root_deleted else root

        helper(root, True)
        return res
複製代碼

By the way，做爲一個 JavaScript 寫手，初看 python 以爲好有趣哦～😆