👨‍💻 738LeetCode

Diagonal Traverse II

Question

Given a list of lists of integers, nums, return all elements of nums in diagonal order as shown in the below images.

Example

Input:
nums = [
    [ 1, 2, 3 ],
    [ 4, 5, 6 ],
    [ 7, 8, 9 ]
]
Output: [1,4,2,7,5,3,8,6,9]
Input: nums = [
    [1,2,3,4,5],
    [6,7],
    [8],
    [9,10,11],
    [12,13,14,15,16]
]
Output: [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

Solution

  • Failed: Change traversing logic from Diagonal Traverse (backtracking). TLE.
  • Failed: Python built-in zip_longest() method. MLE.
  • Solved by BFS. View this grid as a slanted tree. nums[[0][0] is the root node, nums[1][0] the left child and nums[0][1] the right child. Only count left child at col 0 to avoid double counting.
class Solution:
    def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]:
        from collections import deque
        ret = []
        queue = deque([(0, 0)])
        while queue:
            row, col = queue.popleft()
            ret.append(nums[row][col])
            if col == 0 and row + 1 < len(nums):  # Left child
                queue.append((row + 1, col))
            if col + 1 < len(nums[row]):  # Right child
                queue.append((row, col + 1))
        return ret

Comment

There is a smart solution that generates an index array consisting of cell values, each index the sum of row + col.

Additional useful note: Use 2 stacks to mimic a queue.

stack1 = []
stack2 = []
stack1.append(root)
while stack1 or stack2:
    if stack1:
        root = stack1.pop()
        if root.left:
            stack2.append(root.right)  # mark1
        if root.right:
            stack2.append(root.left)  # mark2
    elif stack2:
        root = stack2.pop()
        if root.left:
            stack1.append(root.right)  # mark1
        if root.right:
            stack1.append(root.left)  # mark2
# This is DFS. To get BFS, simply swap mark1 and mark2
# Can also use built-in dequeue or LifoQueue in Python