498. 对角线遍历

class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        ret = []
        m, n = len(mat), len(mat[0])
        for k in range(m + n - 1):
            if k % 2:
                i, j = max(0, k - n + 1), min(k, n - 1)
                while i < m and j >= 0:
                    ret.append(mat[i][j])
                    i += 1
                    j -= 1
            else:
                i, j = min(k, m - 1), max(0, k - m + 1)
                while i >= 0 and j < n:
                    ret.append(mat[i][j])
                    i -= 1
                    j += 1
        return ret