Problem statement

https://binarysearch.com/problems/Leaf-Pairs-Less-Than-Target-Distance-Away/

Solution

Equal to Leetcode 1530. Number of Good Leaf Nodes Pairs.

Complexity

Time is O(n^2), because we have equation T(n) = 2*T(n//2) + O(n^2). Space can be O(n^2) as well.

Code

class Solution:
    def solve(self, root, dist):
        def dfs(node):
            if not node: return []
            if not node.left and not node.right: return [1]

            left = dfs(node.left)
            right = dfs(node.right)
            self.count += sum(l + r <= dist for l in left for r in right)
            return [n + 1 for n in left + right if n + 1 < dist]
        
        self.count = 0
        dfs(root)
        return self.count