仅使用np.diag()修改矩阵非主对角线的方法咨询

直接用np.diag(A, k)读取非主对角线后修改，不会改变原矩阵——因为np.diag返回的是对应对角线元素的副本，而非视图。要修改原矩阵的非主对角线，得直接定位这些元素的索引并赋值，以下是具体实现：

核心思路 #

利用numpy的高级索引，直接定位非主对角线的行、列位置，再批量赋值。针对不同位置的对角线（主对角线上方/下方），索引的生成逻辑略有不同：

示例代码 #

1. 基础实现（修改主对角线上方第1条对角线，k=1）

import numpy as np

# 创建测试矩阵
A = np.zeros((5, 5))

# 定位主对角线上方第1条对角线的索引
rows = np.arange(A.shape[0] - 1)
cols = rows + 1

# 批量赋值
A[rows, cols] = [1, 2, 3, 4]
print(A)

2. 通用函数（支持任意k值：k>0为上方，k<0为下方，k=0为主对角线）

def set_off_diagonal(matrix, values, k):
    n = matrix.shape[0]
    if k > 0:
        # 主对角线上方第k条对角线
        rows = np.arange(n - k)
        cols = rows + k
    elif k < 0:
        # 主对角线下方第|k|条对角线
        cols = np.arange(n + k)
        rows = cols - k
    else:
        # 主对角线（兼容场景）
        rows, cols = np.diag_indices(n)
    # 赋值，确保values长度匹配对角线元素数量
    matrix[rows, cols] = values

# 测试
A = np.zeros((5, 5))
# 修改下方第1条对角线（k=-1）
set_off_diagonal(A, [10, 20, 30, 40], k=-1)
print(A)

原理说明 #

np.diag提取对角线时返回副本，无法反向修改原矩阵；而直接操作矩阵的索引（A[rows, cols]）是对原矩阵的视图进行修改，所以能直接改变原矩阵的内容，避免循环遍历的低效操作。

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