Diagonal Traverse LeetCode solution in Python

1
class Solution:
2
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
3
        flag, rowNum, colNum = True, len(mat), len(mat[0])
4

5
        total, ans = 0, []
6
        while total <= rowNum + colNum - 2:
7
            iLimited = rowNum - 1 if flag else colNum - 1
8
            jLimited = colNum - 1 if flag else rowNum - 1
9
            i = total if total < iLimited else iLimited
10
            j = total - i
11
            while i >= 0 and j <= jLimited:
12
                if flag:
13
                    ans.append(mat[i][j])
14
                else:
15
                    ans.append(mat[j][i])
16
                i -= 1
17
                j += 1
18
            total += 1
19
            flag = not flag
20
        return ans