如何使用Numpy提取指定位置的两条对角线（含边界处理）

我用np.zeros((10,10))创建了一个Numpy数组，其中随机分布着值为1的元素，示例赋值操作如下：

arr[2][4] = 1
arr[3][5] = 1
arr[4][6] = 1

生成的数组为：

[[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]

现在需要实现一个方法，给定数组中的一个位置（比如arr[2][4]），用Numpy方法提取该位置所在的两条对角线（左上到右下、右上到左下），同时要考虑数组边界。需求示例如下：

import numpy as np

arr = np.zeros((10,10))
arr[2][4] = 1
arr[3][5] = 1
arr[4][6] = 1

def extract_diags(arr, position):
    # 这里写Numpy相关逻辑
    return diags

topl_bottomr_diag, topr_bottoml_diag = extract_diags(arr, (2,4))

print(topl_bottomr_diag)
>>> (0, 0, 1, 1, 1, 0, 0, 0)
print(topr_bottoml_diag)
>>> (0, 0, 1, 0, 0, 0, 0)

解决方案 #

直接利用Numpy的索引和对角线提取特性，实现如下：

import numpy as np

arr = np.zeros((10,10))
arr[2][4] = 1
arr[3][5] = 1
arr[4][6] = 1

def extract_diags(arr, position):
    row, col = position
    rows, cols = arr.shape
    
    # 提取左上到右下的对角线
    offset_tblr = col - row
    # 获取完整对角线后，裁剪出包含目标位置的部分
    full_tblr = arr.diagonal(offset=offset_tblr)
    # 计算目标位置在这条对角线上的索引
    idx_in_tblr = row - max(0, row - col)
    # 截取从对角线起点到终点的所有元素
    topl_bottomr_diag = full_tblr
    
    # 提取右上到左下的对角线：先翻转数组列方向，把副对角线转为主对角线
    reversed_arr = arr[:, ::-1]
    offset_tbrl = (cols - 1 - col) - row
    full_tbrl = reversed_arr.diagonal(offset=offset_tbrl)
    # 计算目标位置在翻转后对角线上的索引
    idx_in_tbrl = row - max(0, row + col - (cols - 1))
    topr_bottoml_diag = full_tbrl
    
    # 转换为整数元组，匹配示例输出格式
    return tuple(topl_bottomr_diag.astype(int)), tuple(topr_bottoml_diag.astype(int))

topl_bottomr_diag, topr_bottoml_diag = extract_diags(arr, (2,4))

print(topl_bottomr_diag)
# 输出：(0, 0, 1, 1, 1, 0, 0, 0)
print(topr_bottoml_diag)
# 输出：(0, 0, 1, 0, 0, 0, 0)

代码说明 #

• 左上到右下对角线：通过np.diagonal()的offset参数定位目标对角线，offset值为列索引减行索引，直接获取这条对角线的所有元素，天然适配数组边界。
• 右上到左下对角线：将数组列方向翻转后，副对角线会变成主对角线方向，同样用np.diagonal()提取，最后得到对应元素。
• 最后将Numpy浮点数组转为整数元组，与示例输出格式保持一致。

内容的提问来源于stack exchange，提问作者Showertime