1110. 删点成林
关键要点
- 通过O(1)时间复杂度确认节点是否需要删除
Set to_deleteSet = new HashSet<>(); Arrays.stream(to_delete).forEach(to_deleteSet::add);
- 使用深度优先搜索(DFS)遍历树
node.left = dfs(node.left, s, ans);
node.right = dfs(node.right, s, ans);
- 当前节点处理逻辑
判断当前节点是否为待删除节点
- 不是则返回节点,保持父子关系;
- 当前节点为待删除节点:
- 添加左孩子、右孩子到结果集中【因为是DFS,所以不用考虑放入结果集中的孩子们是否也应该删除】
- 删除当前节点【return null;】,断绝关系
代码编写
/**
* 1110. 删点成林
* https://leetcode.cn/problems/delete-nodes-and-return-forest/description/
* @param root
* @param to_delete
* @return
*/
public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
if (Objects.isNull(root) || Objects.isNull(to_delete) || to_delete.length == 0) {
return Collections.singletonList(root);
}
// 确认是否需要删除node O(1)
Set<Integer> to_deleteSet = new HashSet<>();
Arrays.stream(to_delete).forEach(to_deleteSet::add);
List<TreeNode> result = new ArrayList<>();
// DFS
if (!Objects.isNull(dfs(root, to_deleteSet, result))) {
result.add(root);
}
return result;
}
/**
* DFS->从下(子节点)向上(父节点)
* @param node
* @param s
* @param ans
* @return
*/
private TreeNode dfs(TreeNode node, Set<Integer> s, List<TreeNode> ans) {
if (Objects.isNull(node)) {
return null;
}
// 左
node.left = dfs(node.left, s, ans);
// 右
node.right = dfs(node.right, s, ans);
if (!s.contains(node.val)) {
// 当前节点不能删除(保持与父节点的链接)
return node;
}
// 当前节点为待删除的节点
if (!Objects.isNull(node.left)) {
ans.add(node.left);
}
if (!Objects.isNull(node.right)) {
ans.add(node.right);
}
// 删除当前节点(断开与父节点的链接)
return null;
}
代码解析
- 如果根节点为空或者删除数组为空,直接返回只含有根节点的列表。
- 将待删除的节点值存入一个哈希集合 to_deleteSet 中。
- 通过深度优先搜索(DFS)遍历树,判断当前节点是否需要删除。如果不需要删除,则递归遍历其左右子节点,并更新左右子节点的链接;如果需要删除,则将其左右子节点加入结果集 ans 中,并将当前节点删除。
- 返回结果集 ans,即删除节点后的根节点集合。