[LeetCode] 1110. Delete Nodes And Return Forest

https://leetcode.com/problems/delete-nodes-and-return-forest/description/

leetcode - Delete Nodes And Return Forest

문제 유형 : 이진트리, 재귀

문제 난이도 : Medium

 

문제

Given the root of a binary tree, each node in the tree has a distinct value.

After deleting all nodes with a value in to_delete, we are left with a forest (a disjoint union of trees).

Return the roots of the trees in the remaining forest. You may return the result in any order.

 

이진트리의 루트가 주어집니다.

트리의 각 노드는 다른 값을 가집니다.

to_delete의 노드들을 지우고, 서로소인 트리들의 집합을 만들어서 forest를 반환하시오.

순서는 상관없습니다.

 

풀이

재귀를 통해 다음을 행한다:

1. 현재 노드가 null이라면, 돌아간다.
2. 만약 현재 값이 to_delete의 값이라면
   1. 부모노드에서 현재 노드로의 자식관계를 끊는다.
   2. 현재노드에서도 양쪽의 자식노드로의 관계를 끊는다.
   3. 왼쪽과 오른쪽 자식을 재귀호출하여 탐색한다.
3. 현재 값이 to_delete의 값이 아니라면
   1. 만약 현재 값의 기존 부모 노드와 연결되어있지 않다면, 새 트리의 root이므로 forest에 넣는다.
   2. 왼쪽과 오른쪽 자식을 재귀호출하여 탐색한다.

처음에는 root의 자식들부터 호출을 시작한다. 그 전에 root가 to_delete에 없으면, root를 forest에 넣는다.

 

to_delete는 set으로 변경하고 삭제 후 set에서 제거시키는 방식으로 최적화시킨다.

 

 

코드

C++

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    void traverse(TreeNode* root, TreeNode* parent, bool isLeftChild, set<int>& to_delete, vector<TreeNode*>& forest) {
        if(!root) return;
        if(to_delete.find(root -> val) != to_delete.end()) {
            to_delete.erase(root -> val);
            if(isLeftChild) {
                parent -> left = NULL;
            } else {
                parent -> right = NULL;
            }
            auto leftChild = root -> left;
            auto rightChild = root -> right;
            root -> left = NULL;
            root -> right = NULL;
            traverse(leftChild, root, true, to_delete, forest);
            traverse(rightChild, root, false, to_delete, forest);
        } else {
            if((isLeftChild && parent -> left == NULL) || (!isLeftChild && parent -> right == NULL)) {
                forest.push_back(root);
            }
            traverse(root -> left, root, true, to_delete, forest);
            traverse(root -> right, root, false, to_delete, forest);
        }
    }

    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        vector<TreeNode*> forest;
        set<int> to_delete_set(to_delete.begin(), to_delete.end());
        if(to_delete_set.find(root -> val) != to_delete_set.end()) {
            to_delete_set.erase(root -> val);
            auto leftChild = root -> left;
            auto rightChild = root -> right;
            root -> left = NULL;
            root -> right = NULL;
            traverse(leftChild, root, true, to_delete_set, forest);
            traverse(rightChild, root, false, to_delete_set, forest);
        } else {
            forest.push_back(root);
            traverse(root -> left, root, true, to_delete_set, forest);
            traverse(root -> right, root, false, to_delete_set, forest);
        }

        return forest;
    }
};