🔗 🟡 1110. Delete Nodes And Return Forest 1511

tags: `Weekly Contest 144` `陣列(Array)` `樹(Tree)` `二元樹(Binary Tree)` `雜湊表(Hash Table)` `深度優先搜尋(DFS)`

題意

給定一個二元樹的根節點 $root$ ，樹中的每個的節點都有唯一的值。

另給定一個陣列 $to\_delete$ ，其中包含一些值，如果節點的值在 $to\_delete$ 中，我們就將該節點從樹中刪除，最後會得到一個森林（一些不相交的樹構成的集合）。

請返回在刪除 $to\_delete$ 中的節點後，森林中的每棵樹的根節點。可以按任意順序返回結果。

思路：深度優先搜尋(DFS)

由於我們需要遍歷整個樹，並根據每個節點是否在 $to\_delete$ 中進行刪除，因此可以使用深度優先搜尋(DFS) 來解決這個問題。

定義一個 DFS 函數 $dfs(root, is\_root)$ ，其中 $root$ 表示當前節點， $is\_root$ 表示當前節點是否是某棵樹的根節點。若當前節點的值在 $to\_delete$ 中，則遞迴處理其左右子樹，並將其左右子樹視為新的根節點（因為當前節點將被刪除）。最後返回 $None$ 以表示刪除了當前節點。若當前節點的值不在 $to\_delete$ 中，則根據 $is\_root$ 判斷是否將當前節點加入答案 $ans$ 中。然後遞迴處理其左右子樹，並更新當前節點的左右子樹。

在定義玩 $dfs$ 函數後，我們從根節點開始調用 $dfs$ 函數，並將 $is\_root$ 設為 $True$ 。最後返回答案 $ans$ 。

此外，為了快速查找 $to\_delete$ 中的值，我們可以使用一個集合(Set) $to\_delete\_set$ 來存儲 $to\_delete$ 中的值。

複雜度分析

時間複雜度： $O(n)$ ，其中 $n$ 是二元樹的節點數，需要遍歷每個節點一次。
空間複雜度： $O(n)$ ，遞迴深度最多為 $n$ ，此外需要額外的空間來存儲 $to\_delete\_set$ 和答案 $ans$ 。

class Solution:
    def delNodes(self, root: Optional[TreeNode], to_delete: List[int]) -> List[TreeNode]:
        ans = []
        to_delete_set = set(to_delete)
        def dfs(root: Optional[TreeNode], is_root: bool) -> Optional[TreeNode]:
            if not root:
                return None
            if root.val in to_delete_set:  # 把左右子樹視為新的根節點
                dfs(root.left, True)  # 遞迴處理左右子樹
                dfs(root.right, True)
                return None  # 刪除自己
            else:
                if is_root: # 如果是根節點，加入答案
                    ans.append(root)  
                root.left = dfs(root.left, False)  # 遞迴處理左右子樹
                root.right = dfs(root.right, False)
                return root
        dfs(root, True)
        return ans

class Solution {
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        vector<TreeNode*> ans;
        unordered_set<int> to_delete_set(to_delete.begin(), to_delete.end());
        function<TreeNode*(TreeNode*, bool)> dfs = [&](TreeNode* root, bool is_root) -> TreeNode* {
            if (!root) return nullptr;
            if (to_delete_set.count(root->val)) {
                dfs(root->left, true);
                dfs(root->right, true);
                return nullptr;
            } else {
                if (is_root) ans.push_back(root);
                root->left = dfs(root->left, false);
                root->right = dfs(root->right, false);
                return root;
            }
        };
        dfs(root, true);
        return ans;
    }
};