코딩테스트/99클럽

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

꼽파 2024. 8. 24. 01:17


/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *     this.val = (val===undefined ? 0 : val)
 *     this.left = (left===undefined ? null : left)
 *     this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @param {number} distance
 * @return {number}
 */
var countPairs = function(root, distance) {
    let result = 0;

    function dfs(node) {
        if (!node) return new Map();
        
        // Leaf node
        if (!node.left && !node.right) {
            let map = new Map();
            map.set(1, 1); // Depth 1 with 1 count
            return map;
        }
        
        // Recursive DFS to get leaf counts from left and right subtrees
        const leftCounts = dfs(node.left);
        const rightCounts = dfs(node.right);
        
        // Count good pairs between leaf nodes in left and right subtrees
        for (let [lDepth, lCount] of leftCounts) {
            for (let [rDepth, rCount] of rightCounts) {
                if (lDepth + rDepth <= distance) {
                    result += lCount * rCount;
                }
            }
        }
        
        // Prepare the count map for the current node
        let currentCounts = new Map();
        for (let [depth, count] of leftCounts) {
            if (depth + 1 <= distance) {
                currentCounts.set(depth + 1, (currentCounts.get(depth + 1) || 0) + count);
            }
        }
        for (let [depth, count] of rightCounts) {
            if (depth + 1 <= distance) {
                currentCounts.set(depth + 1, (currentCounts.get(depth + 1) || 0) + count);
            }
        }
        
        return currentCounts;
    }

    dfs(root);
    return result;
}
728x90