Python 中删除节点并返回森林

假设我们有一棵二叉树的根节点，树中的每个节点都有一个唯一的值。删除所有值为 `to_delete` 中的值的节点后，剩下的节点构成一个森林。我们需要找到这个森林中每棵树的根节点。

如果 `to_delete` 数组为 [3,5]，则输出为：

为了解决这个问题，我们将遵循以下步骤：

  • 定义一个数组 `res`
  • 定义一个方法 `solve()`，它将接收节点、`to_delete` 数组和一个布尔类型的 `info` 参数（表示该节点是否是根节点）。该方法的工作方式如下：
  • 如果节点为空，则返回 null
  • 如果节点的值在 `to_delete` 数组中，则 `flag := true`
  • 如果 `flag` 为 false 且 `is_root` 为 true，则将节点插入 `res`
  • 节点的左子节点 := `solve(节点的左子节点, to_delete, flag)`
  • 节点的右子节点 := `solve(节点的右子节点, to_delete, flag)`
  • 如果设置了 `flag`，则返回 null；否则返回 false
  • 在主方法中调用 `solve()` 方法，例如 `solve(node, to_delete, true)`

让我们来看下面的实现，以便更好地理解：

示例

class Solution(object):
   def delNodes(self, root, to_delete):
      """
      :type root: TreeNode
      :type to_delete: List[int]
      :rtype: List[TreeNode]
      """
      to_delete = set(to_delete)
      self.res = []
      self.solve(root,to_delete,True)
      return self.res
   def solve(self,node,to_delete,is_root):
      if not node:
         return None
      flag = node.val in to_delete
      if not flag and is_root:
         self.res.append(node)
      node.left = self.solve(node.left,to_delete,flag)
      node.right = self.solve(node.right,to_delete,flag)
      return None if flag else node

输入

[1,2,3,4,5,6,7]
[3,5]

输出

[[1,2,null,4],[6],[7]]

更新于：2020年3月5日

187 次浏览

开启您的职业生涯

完成课程获得认证

开始学习
广告