与方阵中第一条对角线平行的对角线

Question

我正在寻找所有与方阵中第一个对角线平行的对角线，使用 python。

    matrix = [[1,2,3,4],
              [5,6,7,8],
              [9,0,1,2],
              [3,4,5,6]]

    for i in range(len(matrix)):
       for j in range(len(matrix)):
           diagonal = []
           if i == j + 1 :
                diagonal.append(mat[i][j])
                print(diagonal, end=" ")

我的第一个对角线：

    [1][6][1][6]

输出是：

    [5] [0] [5]

但我正在寻找他们两个：

    [5] [0] [5] and 
    [9] [4]

Answer 1

如果您安装了numpy，您将免费获得diagonals (with offsets)：

import numpy as np

matrix = np.array( [[1,2,3,4],
              [5,6,7,8],
              [9,0,1,2],
              [3,4,5,6]])

print(matrix.diagonal())          # -> [1 6 1 6]
print(matrix.diagonal(offset=-1)) # -> [5 0 5]
print(matrix.diagonal(offset=-2)) # -> [9 4]

这是一种手工完成的方法：

def diagonal(matrix, offset=0):
    ret = []
    l = len(matrix)
    for i, row in enumerate(matrix):
        j = i+offset
        if j < 0:
            continue
        if j >= l:
            break
        ret.append(row[j])
    return ret

print(diagonal(matrix))            # -> [1, 6, 1, 6]
print(diagonal(matrix, offset=-1)) # -> [5, 0, 5]
print(diagonal(matrix, offset=-2)) # -> [9, 4]

或者更紧凑、更高效的版本，通过列表理解生成对角线：

def diagonal(matrix, offset=0):
    dim = len(matrix)
    return [matrix[i][j] for i, j in
            zip(range(-offset if offset < 0 else 0, dim),
                range(offset if offset > 0 else 0, dim))]