[LeetCode] 2641. Cousins in Binary Tree II

https://leetcode.com/problems/cousins-in-binary-tree-ii/description/?envType=daily-question&envId=2024-10-23

leetcode - Cousins in Binary Tree II

문제 유형: 트리, bfs

문제 난이도: Medium

 

문제

Given the root of a binary tree, replace the value of each node in the tree with the sum of all its cousins' values.

Two nodes of a binary tree are cousins if they have the same depth with different parents.

Return the root of the modified tree.

Note that the depth of a node is the number of edges in the path from the root node to it.

 

이진트리의 루트가 주어진다. 노드들의 값을 사촌들의 값의 합으로 대체하시오.

만약 다른 부모이고 같은 레벨이면 사촌이라 합니다.

 

변형된 트리의 루트를 반환하시오.

 

풀이

두 번의 BFS를 통해서 해결할 수 있다.

첫 번째 BFS: sums배열에 sums[level] = 그 레벨의 노드들의 값의 합.

두 번째 BFS: 자식들의 값을 업데이트한다. 각 자식들의 값 = sums[자식의 level] - 두 자식(또는 한 자식 또는 없을 수도 있음)의 합

 

코드

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:
    TreeNode* replaceValueInTree(TreeNode* root) {
        queue<TreeNode*> q;
        vector<long long> sums;

        q.push(root);
        while(!q.empty()) {
            long long sum = 0;
            int qSize = q.size();
            for(int i = 0; i < qSize; ++i) {
                auto curr = q.front();
                q.pop();
                sum += curr -> val;
                if(curr -> left) q.push(curr -> left);
                if(curr -> right) q.push(curr -> right);
            }
            sums.push_back(sum);
        }

        q.push(root);
        int level = 0;
        while(!q.empty()) {
            int qSize = q.size();
            for(int i = 0; i < qSize; ++i) {
                auto curr = q.front();
                q.pop();

                int childrenSum = 0;
                childrenSum += curr -> left? curr -> left -> val : 0;
                childrenSum += curr -> right? curr -> right -> val : 0;

                if(curr -> left) {
                    curr -> left -> val = sums[level+1] - childrenSum;
                    q.push(curr -> left);
                }
                if(curr -> right) {
                    curr -> right -> val = sums[level+1] - childrenSum;
                    q.push(curr -> right);
                }
            }
            level++;
        }

        root -> val = 0;
        return root;
    }
};