498. Diagonal Traverse (Array, Matrix, Simulation)

링크 - https://leetcode.com/problems/diagonal-traverse/description/?envType=company&envId=facebook&favoriteSlug=facebook-three-months

난이도 - Medium

Intuition

이 문제는 동작 방식을 이해하면 쉽게 풀 수 있다. 먼저 direction을 1 또는 0으로 설정하게 한다(초기값은 1). 그리고 row와 col이 edge를 만날 경우에는 우리는 그것에 대한 경우의 수를 적어 놓아야 한다.

Code

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        ans = []
        rows = len(mat)
        cols = len(mat[0])

        row = col = 0
        direction = 1
        while row < rows and col < cols:
            ans.append(mat[row][col])
            new_row = row + (-1 if direction == 1 else 1)
            new_col = col + (1 if direction == 1 else -1)

            if new_row < 0 or new_col < 0 or new_row >= rows or new_col >= cols:
                if direction == 1:
                    row += (col == cols-1)
                    col += (col < cols-1)
                else:
                    col += (row == rows-1)
                    row += (row < rows-1)
                direction = 1 - direction
            else:
                row = new_row
                col = new_col

        return ans

Complexity

Time Complexity: O(N*M)

Space Complexity: O(1)