498. Diagonal Traverse - cocoder39/coco39_LC GitHub Wiki

498. Diagonal Traverse

control moving direction (ie up_right) to simulate the process.

• Boundary Conditions:
  • When moving up-right, if you exceed the upper or right boundary:
    • If you can move right (not in the last column), move to (row, col + 1).
    • Otherwise, move down to the next row (row + 1, col).
  • When moving down-left, if you exceed the lower or left boundary:
    • If you can move down (not in the last row), move to (row + 1, col).
    • Otherwise, move right to the next column (row, col + 1).

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        if not mat or not mat[0]:
            return []
        
        row, col = len(mat), len(mat[0])
        r, c = 0, 0
        up_right = True
        res = []
        while len(res) < row * col:
            res.append(mat[r][c])
            new_row = r + (-1 if up_right else 1)
            new_col = c + (1 if up_right else -1)

            if 0 <= new_row < row and 0 <= new_col < col:
                r, c = new_row, new_col
            else:
                if up_right:
                    # cannot swap the if else order
                    # once reach upper right corner mat[0][col-1], can only go down
                    # if swap order, once reach upper right corner, it will go right 
                    if c == col - 1:
                        r += 1
                    else: # r == 0
                        c += 1
                else:
                    if r == row - 1:
                        c += 1
                    else: # c == 0
                        r += 1
                up_right = not up_right
        return res

option 2: group by diagonal similar to leetcode 1424

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        if not mat or not mat[0]:
            return []
        
        row, col = len(mat), len(mat[0])
        group = collections.defaultdict(list)
        
        for r in range(row-1, -1, -1):
            for c in range(col):
                diagonal = r + c
                group[diagonal].append(mat[r][c])
        
        up_right = True
        diagonal = 0
        res = []
        while diagonal in group:
            if up_right:
                res.extend(group[diagonal])
            else:
                res.extend(group[diagonal][::-1])
            diagonal += 1
            up_right = not up_right
        return res