如何用Python生成含特定对角线结构的矩阵？

你当前的代码会把所有索引差1的位置都设为1，但忽略了关键情况：当两个索引属于原N×N网格不同行的首尾时（比如N=3时，索引2和3），这两个位置在物理上并不相邻，因此应该保持0而非设为1。

修正方案1：改进循环逻辑 #

通过取模运算判断当前索引是否处于行首或行尾，仅对同一行内的相邻索引赋值1：

import numpy as np

N = 3
D2 = np.zeros((N**2, N**2))

for i in range(N**2):
    # 主对角线元素设为-2
    D2[i, i] = -2
    # 处理右侧相邻元素（j = i+1）
    if i < N**2 - 1:
        # 若当前索引不是行的最后一个元素，才赋值1
        if (i + 1) % N != 0:
            D2[i, i+1] = 1
    # 处理左侧相邻元素（j = i-1）
    if i > 0:
        # 若当前索引不是行的第一个元素，才赋值1
        if i % N != 0:
            D2[i, i-1] = 1

修正方案2：用numpy对角线构造（更高效） #

避免双重循环，直接利用np.diag构造不同位置的对角线，适合大N场景：

import numpy as np

N = 3
size = N**2

# 初始化主对角线为-2的矩阵
D2 = np.diag([-2] * size)

# 构造上对角线（k=1，对应j = i+1）：仅当i+1不是行尾时设为1
upper_diag = np.where(np.arange(1, size) % N != 0, 1, 0)
D2 += np.diag(upper_diag, k=1)

# 构造下对角线（k=-1，对应j = i-1）：仅当i不是行首时设为1
lower_diag = np.where(np.arange(size-1) % N != N-1, 1, 0)
D2 += np.diag(lower_diag, k=-1)

两种方案生成的矩阵都是N²×N²的分块对角矩阵，每个子块是N×N的三对角矩阵，主对角线为-2，同一子块内的上下次对角线为1，子块间元素全为0。

内容的提问来源于stack exchange，提问作者Jennifer Cuellar Castrillon