# 1530. Number of Good Leaf Nodes Pairs

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
"""
     1
 2      3
4 5    6

depths_of_all_child_leaves(node_4) is [0]
depths_of_all_child_leaves(node_5) is [0]
depths_of_all_child_leaves(node_2) is [1, 1] named as A
A_1 is [1+1, 1+1]

depths_of_all_child_leaves(node_6) is [0]
depths_of_all_child_leaves(node_3) is [1] named as B
B_1 is [1+1]

depths_of_all_child_leaves(node_1) is A_1 + B_1

"""
class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        self.good_pairs = 0  # be global to all dfs function-calls

        def depths_of_all_child_leaves(node):
            """Return a list of depths of all its child-leaves, could be itself."""
            if not node:
                return []
            if not node.left and not node.right:  # termination
                return [0]
            
            left_depths = depths_of_all_child_leaves(node.left)
            right_depths =  depths_of_all_child_leaves(node.right)
            
            left_depths_add_1 = [x + 1 for x in left_depths]
            right_depths_add_1 = [x + 1 for x in right_depths]
            
            # all lengths between: left-child leaves and right-child leaves
            # and pass this `node`; these will be shortest
            len_shortest_paths_thru_me = [
                x + y for x in left_depths_add_1 for y in right_depths_add_1]
            self.good_pairs += sum(length <= distance for length in len_shortest_paths_thru_me)
            
            return left_depths_add_1 + right_depths_add_1
            
        depths_of_all_child_leaves(root)  # time: O(N^3); N is num of nodes
        return self.good_pairs
    
class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        self.good_pairs = 0  # be global to all dfs function-calls

        def depths_of_all_child_leaves(node):
            """Return a list of depths of all its child-leaves, could be itself."""
            if not node:
                return []
            if not node.left and not node.right:  # termination
                return [0]
            
            left_depths = depths_of_all_child_leaves(node.left)
            right_depths =  depths_of_all_child_leaves(node.right)
            
            left_depths_add_1 = [x + 1 for x in left_depths]
            right_depths_add_1 = [x + 1 for x in right_depths]
            
            # all lengths between: left-child leaves and right-child leaves
            # and pass this `node`; these will be shortest
            len_shortest_paths_thru_me = [
                x + y for x in left_depths_add_1 for y in right_depths_add_1]
            self.good_pairs += sum(length <= distance for length in len_shortest_paths_thru_me)
            
            # remove those already too large; just for optimization
            return [x for x in left_depths_add_1 + right_depths_add_1 if x <= distance]
            
        depths_of_all_child_leaves(root)  # time: O(N * 2^distance * 2^distance)
        return self.good_pairs
    
class Solution:
    def countPairs(self, root: TreeNode, distance: int) -> int:
        """
        Use counter instead of all literal int.
        """
        self.good_pairs = 0  # be global to all dfs function-calls

        def depths_of_all_child_leaves(node):
            """Return a counter for depths of all its child-leaves, could be itself."""
            counter = [0] * (distance + 1)  # handmade counter for 0 ~ distance
            if not node:
                return counter
            if not node.left and not node.right:  # termination
                counter[0] = 1
                return counter
            
            left_depths = depths_of_all_child_leaves(node.left)
            right_depths =  depths_of_all_child_leaves(node.right)
            
            # everyone plus 1, so counter if shifted; 
            left_depths[1::] = left_depths[0:-1]
            right_depths[1::] = right_depths[0:-1]
            left_depths[0] = 0
            right_depths[0] = 0
            # If you want to use for-loop,
            # note that we have to update from the rear;
            # otherwise, li[0] will overwrite every slot
                         
            # all lengths between: left-child leaves and right-child leaves
            # and pass this `node`; these will be shortest
            for x, ct_x in enumerate(left_depths):
                for y, ct_y in enumerate(right_depths):
                    if x + y <= distance:
                        self.good_pairs += ct_x * ct_y
            
            # combine
            return [x + y for x, y in zip(left_depths, right_depths)]

            
        depths_of_all_child_leaves(root)  # time: O(N * distance^2)
        return self.good_pairs