BinarySearch 0687 Counting K-Length Paths on Binary Tree

Problem statement

https://binarysearch.com/problems/Counting-K-Length-Paths-on-Binary-Tree/

Solution

The idea is for each node keep array of length k + 1: how many children we have for each depth.

Complexity

It is O(nk) for time and O(hk) for space.

Code

class Solution:
    def solve(self, root, k):
        self.ans = 0
        def dfs(node):
            if not node: return [0] * (k + 1)
            lft = dfs(node.left)
            rgh = dfs(node.right)
            self.ans += sum(lft[x]*rgh[k-3-x] for x in range(k-1))
            ans = [1] + [0] * k
            for i in range(1, k+1):
                ans[i] += lft[i-1] + rgh[i-1]
            self.ans += ans[k-1]
            return ans

        dfs(root)
        return self.ans