BinarySearch 0557 Diagonal Tree Traversal

Problem statement

https://binarysearch.com/problems/Diagonal-Tree-Traversal/

Solution

The idea is to traverse tree, where we keep coordinates of nodes.

Complexity

Time complexity is O(n), space complexity is O(h)

Code

class Solution:
    def solve(self, root):
        cnt = Counter()
        def dfs(node, i, j):
            cnt[i - j] += node.val
            if node.left: dfs(node.left, i + 1, j - 1)
            if node.right: dfs(node.right, i + 1, j + 1)

        dfs(root, 0, 0)
        t = len(cnt)
        return [cnt[2*i] for i in range(t)]