Computer >> Hướng Dẫn Máy Tính >  >> Lập Trình >> Python

Delete Nodes And Return Forest: Xóa Node Trả Về Forest (Danh Sách Cây Con)

Đề bài: Cho root cây nhị phân và mảng to_delete. Xóa tất cả node có giá trị trong to_delete. Trả về list root của các cây còn lại (forest).

Ví dụ: Tree [1,2,3,4,5,6,7], to_delete=[3,5] → Forest roots: [1, 6, 7] (cây 1-2-4, cây 6, cây 7).

Ý tưởng: Post-order DFS + Flag is_root

Duyệt post-order. Khi xóa node, children của nó trở thành root mới (nếu không bị xóa).

  • Hàm dfs(node, is_root) trả về node mới (hoặc None nếu bị xóa).
  • Nếu node.val trong to_delete:
    • Đệ quy children với is_root=True.
    • Trả về None (node này bị xóa).
  • Ngược lại:
    • Nếu is_root=True → thêm vào result.
    • Đệ quy children với is_root=False.
    • Trả về node.

Code Python


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def delNodes(self, root: TreeNode, to_delete: list[int]) -> list[TreeNode]:
        to_delete_set = set(to_delete)
        self.forest = []
        
        def dfs(node, is_root):
            if not node:
                return None
            
            deleted = node.val in to_delete_set
            
            if not deleted and is_root:
                self.forest.append(node)
            
            # Recurse: children become roots if current is deleted
            node.left = dfs(node.left, deleted)
            node.right = dfs(node.right, deleted)
            
            return None if deleted else node
        
        dfs(root, True)
        return self.forest

# Helper: build tree & print forest
def build_tree(vals):
    if not vals: return None
    nodes = [TreeNode(v) if v is not None else None for v in vals]
    kids = nodes[::-1]
    root = kids.pop()
    for node in nodes:
        if node:
            if kids: node.left = kids.pop()
            if kids: node.right = kids.pop()
    return root

def inorder(root):
    return inorder(root.left) + [root.val] + inorder(root.right) if root else []

# Test
root = build_tree([1,2,3,4,5,6,7])
ob = Solution()
forest = ob.delNodes(root, [3,5])
for i, tree in enumerate(forest):
    print(f"Tree {i}: {inorder(tree)}")
# Tree 0: [4, 2, 1]
# Tree 1: [6]
# Tree 2: [7]

Phân tích

  • Time: O(n) - duyệt mỗi node 1 lần.
  • Space: O(h + d) recursion stack + to_delete set (d = len(to_delete)).
  • Post-order đảm bảo children được xử lý trước parent.

Kết luận

DFS post-order với flag is_root là pattern chuẩn cho "delete nodes return forest". Xử lý tách cây con thành tree mới khi parent bị xóa.