python实现二维数组对角线遍历

一、题目要求

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

示例：

输入:

[

 [ 1, 2, 3 ],

 [ 4, 5, 6 ],

 [ 7, 8, 9 ]

]

输出:  [1,2,4,7,5,3,6,8,9]

二、思路

1 遍历可以理解为遍历对角线，对角线的特性是 行索引 i + 列索引 j = 对角线axis索引

2 假设axis为偶数从下网上遍历，axis为奇数从上到下

由于遍历时，奇数我从行索引为0开始遍历，实现过程中遇到一些小问题（这里以奇数为例）：

下图蓝箭头的遍历行索引会从0开始，但是这里列索引已经溢出了，需要加一个列索引溢出判断就能解决（i-j就是指对应的列索引，continue后行会+1，列就会减1）。

# 防止行索引(i-j)溢出
if i - j > m - 1:
    continue

三、代码

#!/usr/bin/python3
# author lzx

# python实现
# 思路：按斜轴数遍历，偶数从下到上，奇数从上到下
import numpy as np
def func(m, n):
    list1 = np.arange(1, m * n + 1).reshape(m, n).tolist()
    print(list1)
    target = []
    axis = m + n - 1

    for i in range(axis):
        # 偶数从下往上遍历
        if i % 2 == 0:
            for j in range(n):
                # 行索引(i-j)不能为负数
                if (i - j) < 0:
                    break
                # 防止行索引(i-j)溢出
                if i - j > m - 1:
                    continue
                target.append(list1[i - j][j])

        # 奇数从上往下遍历
        else:
            for k in range(m):
                if (i - k) < 0:
                    break
                if i - k > n - 1:
                    continue
                target.append(list1[k][i - k])
    return target


if __name__ == '__main__':
    m = int(input('输入m：'))
    n = int(input('输入n：'))
    target = func(m, n)
    print(target)