如何根据指定索引获取矩阵向下延伸的左右两条对角线元素？

定义说明 #

对于矩阵中坐标为(i,j)的元素（i为行索引，j为列索引）：

• 向下右对角线：从该位置出发，每步进位规则为「行号+1、列号+1」，直到超出矩阵边界的所有元素集合
• 向下左对角线：从该位置出发，每步进位规则为「行号+1、列号-1」，直到超出矩阵边界的所有元素集合

实现代码 #

import numpy as np

def get_lower_diags(matrix, row_idx, col_idx):
    max_row, max_col = matrix.shape
    # 取向下右对角线
    diag_right = []
    cur_r, cur_c = row_idx + 1, col_idx + 1
    while cur_r < max_row and cur_c < max_col:
        diag_right.append(matrix[cur_r][cur_c])
        cur_r += 1
        cur_c += 1
    # 取向下左对角线
    diag_left = []
    cur_r, cur_c = row_idx + 1, col_idx - 1
    while cur_r < max_row and cur_c >= 0:
        diag_left.append(matrix[cur_r][cur_c])
        cur_r += 1
        cur_c -= 1
    return diag_right, diag_left

测试效果 #

用示例矩阵测试：

matrix = np.array([[2, 0, 0, 2],
                  [3, 9, 8, 3],
                  [3, 0, 0, 2],
                  [0, 0, 0, 0]])

# 测试坐标(1,2)
diag_right, diag_left = get_lower_diags(matrix, 1, 2)
print(f"当前元素值：{matrix[1][2]}")
print(f"diag_right = {diag_right}")
print(f"diag_left = {diag_left}")

输出结果完全匹配预期：

当前元素值：8
diag_right = [2]
diag_left = [0, 0]

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