对角线遍历

题目

给定一个含有M×N个元素的矩阵（M行，N列），请以对角线遍历的顺序返回这个矩阵中的所有元素，对角线遍历如下图所示。

思路

先简化问题，一步一步地来

1. 确定循环次数

2. 方向变更

方向跟遍历次数的奇偶性相关，因此遍历次数对2取余即可。
大致框架

m, n = len(matrix), len(matrix[0])
count = m + n - 1
answer = [] # 用于保存遍历结果
start_point = (0, 0)
for i in range(count):
    if i % 2 == 0:
        # 下一步遍历起始点为start_point的右上方向的对角线元素
        pass
    else:
        # 下一步遍历起始点为start_point的右上方向的对角线元素
        pass
return answer

4. 遍历起始点为（x，y）的右上角对角线元素

右上方向移动：x -= 1, y += 1
边界条件 x >= 0, y < n (横坐标不为负数，纵坐标不大于列数）

考虑越界处理

while x >= 0 and y < n:
    answer.append(matrix[x][y])
    x -= 1
    y += 1
if y < n:
    x += 1
else:
    x += 2
    y -= 1

5.遍历起始点为（x, y)的左下角对角线元素

左下方向移动：x += 1 ， y -= 1
边界条件：x < m , y >= 0

while x < m and y >= 0:
    answer.append(matrix[x][y])
    x += 1
    y -= 1
if x < m:
    y += 1
else:
    x -= 1
    y += 2

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        if len(matrix) == 0: return []
        m, n = len(matrix), len(matrix[0])
        count = m + n - 1
        answer = [] # 用于保存遍历结果
        x, y = 0, 0 # start_point = (0, 0)
        for i in range(count):
            if i % 2 == 0:
                # 下一步遍历起始点为start_point的右上方向的对角线元素
                while x >= 0 and y < n:
                    answer.append(matrix[x][y])
                    x -= 1
                    y += 1
                if y < n:
                    x += 1
                else:
                    x += 2
                    y -= 1
            else:
                # 下一步遍历起始点为start_point的右上方向的对角线元素
                while x < m and y >= 0:
                    answer.append(matrix[x][y])
                    x += 1
                    y -= 1
                if x < m:
                    y += 1
                else:
                    x -= 1
                    y += 2
        return answer