[LeetCode] 1530. Number of Good Leaf Nodes Pairs

https://leetcode.com/problems/number-of-good-leaf-nodes-pairs/description/

leetcode - Number of Good Leaf Nodes Paris

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

문제 난이도 : Medium

 

문제

You are given the root of a binary tree and an integer distance. A pair of two different leaf nodes of a binary tree is said to be good if the length of the shortest path between them is less than or equal to distance.

Return the number of good leaf node pairs in the tree.

 

이진트리의 root를 받는다. 리프 노드 쌍의 거리가 distance이하인 경우 좋은 쌍이라고 한다.

좋은 쌍의 개수를 구하시오.

 

풀이

pairs를 0으로 초기화한다.

 

재귀적으로 다음의 함수를 수행한다:

만약 현재 노드가 null이면, {}을 반환한다.

만약 현재 노드가 리프 노드이면, 즉 양쪽 자식이 모두 없으면, {1}을 반환한다.

 

리프 노드가 아니라면, 양쪽의 자식으로부터 리프까지의 거리가 담긴 배열을 재귀 호출을 통해 가져온다.

두 양쪽 자식배열들로 조합을 만들어서 각 조합의 합이 distance이하인 경우 pairs를 1씩 증가시킨다.

 

두 배열을 합치고 각 요소에 1을 더해서 부모 노드에게 반환한다. 단, 1을 더했을 때 distance보다 크거나 같다면, 그 요소는 제외해야 한다.

 

pairs를 반환한다.

 

코드

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 {
private:
    int pairs;
public:
    vector<int> traverse(TreeNode* root, int distance) {
        if(!root) return {};
        if(!(root -> left) && !(root -> right)) return {1};
        vector<int> leftLeaves = traverse(root -> left, distance);
        vector<int> rightLeaves = traverse(root -> right, distance);

        for(const auto &leftLeaf : leftLeaves) {
            for(const auto &rightLeaf : rightLeaves) {
                if(leftLeaf + rightLeaf <= distance) pairs++;
            }
        }

        vector<int> leaves;
        for(const auto &leftLeaf : leftLeaves) {
            if(leftLeaf+1 >= distance) continue;
            leaves.push_back(leftLeaf+1);
        }
        for(const auto &rightLeaf : rightLeaves) {
            if(rightLeaf+1 >= distance) continue;
            leaves.push_back(rightLeaf+1);
        }

        return leaves;
    }
    int countPairs(TreeNode* root, int distance) {
        pairs = 0;
        traverse(root, distance);
        return pairs;
    }
};