使用 numpy 在矩阵中设置交替的非对角线元素

Question

我想在任意 NxN 矩阵中设置交替的非对角线元素，如图所示。

目前我使用的是蛮力方式，不适应矩阵大小的变化：

import numpy as np
dim = 8
size=dim**2
v=0.1
w=1

h = np.zeros(size)

valuesv = [1,8,19,26,37,44,55,62]
valuesw = [10,17,28,35,46,53]

h[valuesv] = v
h[valuesw] = w

ham=h.reshape((int(np.sqrt(size)),int(np.sqrt(size))))

有没有更优雅的方式来填充矩阵？例如。使用 numpy 的 fill_diagonal？

Answer 1

下面的呢？我的回答是基于this

import numpy as np

# array size
nn = 8

aa = np.zeros((nn,nn),dtype='float64')

# get indices
rows, cols = np.indices((nn,nn))

# get indices of the fist upper minor diagonal
#      k = -1 will get you the first lower minor diagonal
#      k = 2  will get you the second upper minor diagonal
row_upper = np.diag(rows, k=1)
col_upper = np.diag(cols, k=1)

# create data for the upper minor diagonal
upper_data = np.zeros(len(row_upper))
# create alternating data
upper_data[0::2] = 0.1
upper_data[1::2] = 1.0

# set data
aa[row_upper,col_upper] = upper_data

编辑

使用更少内存的解决方案。无需使用np.indices 创建索引矩阵。

import numpy as np
# array size
nn = 10
aa = np.zeros((nn,nn),dtype='float64')
aa.ravel()[1::(2*nn+2)]=0.1          # set upper diagonal 0.1s
aa.ravel()[2+nn::(2*nn+2)]=1.0       # set upper diagonal 1.0s
aa.ravel()[nn::(2*nn+2)]=0.1         # set lower diagonal 0.1s
aa.ravel()[2*nn+1::(2*nn+2)]=1.0     # set lower diagonal 1.0s