[LeetCode] 2458. Height of Binary Tree After Subtree Removal Queries

https://leetcode.com/problems/height-of-binary-tree-after-subtree-removal-queries/description/?envType=daily-question&envId=2024-10-26

leetcode - Height of Binary Tree After Subtree Removal Queries

문제 유형: 재귀, dfs, 다이나믹 프로그래밍

문제 난이도: Hard

 

문제

You are given the root of a binary tree with n nodes. Each node is assigned a unique value from 1 to n. You are also given an array queries of size m.

You have to perform m independent queries on the tree where in the ith query you do the following:

• Remove the subtree rooted at the node with the value queries[i] from the tree. It is guaranteed that queries[i] will not be equal to the value of the root.

Return an array answer of size m where answer[i] is the height of the tree after performing the ith query.

Note:

• The queries are independent, so the tree returns to its initial state after each query.
• The height of a tree is the number of edges in the longest simple path from the root to some node in the tree.

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

각 노드는 1부터 n까지의 고유한 번호를 가진다.

queries라는 배열을 받는다, 요소는 m개있다.

 

각 쿼리마다 이렇게 처리하라:

queries[i]번의 키를 가진 노드를 루트로 하는 서브트리를 없앴다고 가정하고, 트리의 높이를 구하시오.

그 답안은 배열에 담아서 반환하시오.

 

풀이

입력 크기를 보면 알겠듯이, 매번 없다고 가정하고 루트부터 트리를 순회하면 TLE에 걸린다.

height[i]에는 i번째 노드의 높이를 저장한다.

removal[i]에는 i를 키로하는 루트부터 없앴다고 가정하고 만든 트리의 높이이다.

 

dfs를 통해서 removal[i]를 채울 것이다.

 

현재 노드의 깊이에는 maxLevel을 저장한다.

그리고, 왼쪽 자식에 대해 호출하는데, level은 1만 높이면 되고, maxLevel은 maxLevel과 1 + level + h(오른쪽자식)으로 넣으면 된다.

이게 무슨 의미냐면, 만약 왼쪽 자식이 없어지고, 오른쪽 형제 노드에 더 긴 브랜치가 있어서 최대 깊이가 따로 있다면, 그 값을 다음 재귀에서 저장하는데 쓰라는 뜻이다.

 

순수 깊이만을 구하는 h에서는

만약 현재 루트가 null이면, -1을 반환한다. 이러면 1이 추가로 더해진 것을 방지할 수 있다.

그게 아니면, 1 + max(h(왼쪽자식), h(오른쪽자식))를 구해서 height에 저장한다.

 

이렇게 dfs를 통해서 removal 테이블을 정리해놓고, 각 쿼리에 대한 값을 removal에서 가져와서 배열에 담아 저장하고 반환하면 된다.

 

코드

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:
    int height[100001], removal[100001];
    int h(TreeNode* root) {
        if(!root) return -1;
        int x = root -> val;
        if(height[x] != -1) return height[x];
        return height[x] = 1+max(h(root -> left), h(root -> right));
    }
    void dfs(TreeNode* root, int level, int maxLevel) {
        if(!root) return;
        removal[root -> val] = maxLevel;
        dfs(root -> left, level+1, max(maxLevel, 1 + level + h(root -> right)));
        dfs(root -> right, level+1, max(maxLevel, 1 + level + h(root -> left)));
    }
    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        memset(height, -1, sizeof(height));
        memset(removal, -1, sizeof(removal));
        dfs(root, 0, 0);

        vector<int> ans;
        for(const auto &query : queries) {
            ans.push_back(removal[query]);
        }
        return ans;
    }
};