如何基于自定义逻辑利用指定列表创建二维NumPy数组？

给定数值列表：

k_list = [5,6,4,7,8]

需要基于该列表构建一个如下结构的NumPy二维数组：

arr = np.array([[k1+k2, -k2, 0, 0, 0],
                [-k2, k2+k3, -k3, 0, 0],
                [0, -k3, k3+k4, -k4, 0],
                [0, 0, -k4, k4+k5, -k5],
                [0, 0, 0, -k5, k5]])

其中k1到k5对应列表中的元素，预期输出为：

arr = np.array([[11, -6, 0, 0, 0],
                [-6, 10, -4, 0, 0],
                [0, -4, 11, -7, 0],
                [0, 0, -7, 15, -8],
                [0, 0, 0, -8, 8]])

解决方案 #

利用NumPy的数组切片和对角矩阵特性，可高效构建目标矩阵，代码如下：

import numpy as np

k_list = [5,6,4,7,8]
n = len(k_list)
# 创建n×n的零矩阵作为基础
arr = np.zeros((n, n), dtype=int)

# 填充主对角线元素
main_diag = np.concatenate([np.add(k_list[:-1], k_list[1:]), [k_list[-1]]])
arr[np.arange(n), np.arange(n)] = main_diag

# 填充主对角线上下两侧的次对角线元素
off_diag = -np.array(k_list[1:])
arr[np.arange(n-1), np.arange(1, n)] = off_diag  # 上侧次对角线
arr[np.arange(1, n), np.arange(n-1)] = off_diag  # 下侧次对角线

print(arr)

代码说明 #

• 先创建对应尺寸的零矩阵作为基础框架；
• 主对角线：前n-1个元素为列表中相邻元素的和，最后一个元素直接取列表末尾值；
• 上下次对角线：均取列表第二个元素开始的负值，分别填充到主对角线的右上方和左下方位置；

运行代码后输出结果与预期一致：

[[11 -6  0  0  0]
 [-6 10 -4  0  0]
 [ 0 -4 11 -7  0]
 [ 0  0 -7 15 -8]
 [ 0  0  0 -8  8]]

该方法具备可扩展性，无论输入列表长度如何，都能自动生成对应结构的矩阵，无需手动调整代码框架。

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