[LeetCode] 2265. Count Nodes Equal to Average of Subtree

https://leetcode.com/problems/count-nodes-equal-to-average-of-subtree/description/

Count Nodes Equal to Average of Subtree - LeetCode

leetcode - Count Nodes Equal to Average of Subtree

문제 유형 : 트리, 재귀, 순회, DFS

문제 난이도 : Medium

 

문제

Given the root of a binary tree, return the number of nodes where the value of the node is equal to the average of the values in its subtree.

Note:

• The average of n elements is the sum of the n elements divided by n and rounded down to the nearest integer.
• A subtree of root is a tree consisting of root and all of its descendants.

이진 트리의 루트가 주어집니다. 노드의 값이 자신을 루트로 하는 서브트리의 평균과 같은 노드들의 개수를 구하시오.

 

n개의 요소의 평균의 소수점은 잘립니다.

루트의 서브트리는 루트와 자식들을 포함하는 트리를 말합니다.

 

 

풀이

후위 순회를 통해서 양쪽 자식들이 각 루트가 되는 서브트리의 합과 개수를 불러와서 계산할 수 있다.

 

후위 순회 함수가 <sum, num>(합, 개수)를 반환하도록 하게 한다.

만약 루트가 nil이면, <0, 0>을 반환한다.

양쪽 서브트리와 루트의 값의 합을 구한 sum을 재귀 호출을 통해서 구한다.

총 개수도 함께 구한다.

평균을 합 / 개수로 구한다. 나머지는 버리면 된다.

평균이 현재 루트의 값과 같다면, 카운트를 1 증가시킨다.

<합, 개수>를 반환시킨다.

 

 

코드

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 {
    int count = 0;
public:
    pair<int, int> postOrder(TreeNode* rooot) {
        if(rooot == nullptr) return make_pair(0, 0);

        auto left = postOrder(rooot -> left);
        auto right = postOrder(rooot -> right);
        int sum = left.first + right.first + rooot -> val;
        int num = left.second + right.second + 1;
        if(sum / num == rooot -> val) {
            count++;
        }
        return make_pair(sum, num);
    }

    int averageOfSubtree(TreeNode* root) {
        ios_base::sync_with_stdio(0);
        cin.tie(0);
        cout.tie(0);
        
        auto x = postOrder(root);

        return count;
    }
};