每日算法题—删点成林

题目描述

给出二叉树的根节点 root，树上每个节点都有一个不同的值。

如果节点值在 to_delete 中出现，我们就把该节点从树上删去，最后得到一个森林（一些不相交的树构成的集合）。

返回森林中的每棵树。
来源：https://leetcode-cn.com/problems/delete-nodes-and-return-forest

如输入一棵树和待删除集合[3,5]：

返回：[[1,2,null,4],[6],[7]]，即树因为被删了3和5两个节点，所以变成了3棵树，分别是[1,2,null,4]、[6]、[7]

思路

使用深度优先遍历，当遇到待删除节点时，如果该节点有孩子节点，那么孩子节点会独立成树，同时还需要将该节点和其父节点与左右子节点断开

实现

 fun delNodes(root: TreeNode?, to_delete: IntArray): List<TreeNode?> {
        if (root == null) {
            return arrayListOf()
        }
        val forest = arrayListOf(root)
        reverse(root, null, true, to_delete, forest)//根节点没有父节点，所以传null
        return forest
    }

    //isLeft表示node是parent的左节点还是右节点，用于parent需要删除node时确定是删左孩子还是右孩子
    fun reverse(node: TreeNode, parent: TreeNode?, isLeft: Boolean, to_delete: IntArray, forest: ArrayList<TreeNode>) {
        node.left?.let {
            reverse(it, node, true, to_delete, forest)
        }
        node.right?.let {
            reverse(it, node, false, to_delete, forest)
        }
        if (to_delete.contains(node.value)) {//该节点是待删除节点
            node.left?.let { forest.add(it) }//左孩子不为空，则左孩子会独立成树
            node.right?.let { forest.add(it) }//右孩子同理
            parent?.let { if (isLeft) parent.left = null else parent.right = null }//删除该节点和其父节点的关联
            node.left = null//删除该节点和其子节点的关联
            node.right = null
            forest.remove(node)//考虑删除根节点时，因为根节点一开始就在forest中，所以需要把它从forest中删除
        }
    }