99클럽 코테 스터디 33일차 TIL | Number of Good Leaf Nodes Pairs

🖥️ 문제

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.

Example 1:

Input: root = [1,2,3,null,4], distance = 3
Output: 1
Explanation: The leaf nodes of the tree are 3 and 4 and the length of the shortest path between them is 3. This is the only good pair.

📝 풀이

class Solution {
    public int countPairs(TreeNode root, int distance) {
        int[] result = new int[1];
        dfs(root, distance, result);
        return result[0];
    }
    
    private int[] dfs(TreeNode node, int distance, int[] result) {
        if (node == null) return new int[distance + 1];
        
        if (node.left == null && node.right == null) {
            int[] leafDist = new int[distance + 1];
            leafDist[1] = 1;
            return leafDist;
        }
        
        int[] leftDist = dfs(node.left, distance, result);
        int[] rightDist = dfs(node.right, distance, result);
        
        for (int i = 1; i <= distance; i++) {
            for (int j = 1; j + i <= distance; j++) {
                result[0] += leftDist[i] * rightDist[j];
            }
        }
        
        int[] currentDist = new int[distance + 1];
        for (int i = 1; i < distance; i++) {
            currentDist[i + 1] = leftDist[i] + rightDist[i];
        }
        
        return currentDist;
    }
}

문제는 주어진 이진 트리에서 좋은 리프 노드(distance 이하일 때) 쌍의 개수를 세는 문제다. 먼저, 트리의 각 노드를 순회하면서, 각 노드에서 리프 노드까지의 거리를 계산했다. 이때, 노드가 리프 노드라면 해당 노드로부터의 거리를 찾았다. 이후 각 노드에서 왼쪽 서브트리와 오른쪽 서브트리의 리프 노드 간 거리를 조합하여 좋은 리프 노드 쌍의 개수를 계산했다.