使用numpy as_strided检索以主对角线为中心的子数组

import numpy as np

def get_diagonal_subarray(arr, k):
    shape = arr.shape
    strides = arr.strides
    subarray_shape = (min(shape[0], shape[1] - abs(k)),)
    subarray_strides = (strides[0] + strides[1] * k,)
    subarray = np.lib.stride_tricks.as_strided(arr, shape=subarray_shape, strides=subarray_strides)
    return subarray

# 示例输入矩阵
matrix = np.array([[1, 2, 3, 4],
                   [5, 6, 7, 8],
                   [9, 10, 11, 12]])

# 检索以主对角线为中心的子数组
subarray = get_diagonal_subarray(matrix, 0)
print(subarray)

[1 6 11]