尝试找到Python式的方法来部分填充numpy数组

Question

我有一个形状为 (3,1000) 的 numpy 数组

psi

psi.__class__
Out[115]: numpy.ndarray

psi.shape
Out[116]: (3, 1000)

我想用另一个数组部分填充

psi

b

b.__class__
Out[113]: numpy.ndarray

b.shape
Out[114]: (3, 500)

我可以用循环来做到这一点：

for n in range(3):
   psi[n][:500] = b[n]

但在我看来，应该有一种更直接的方法来做到这一点。但举个例子

psi[:][:500] = b

因错误而失败

Traceback (most recent call last):

  File "<ipython-input-120-6b23082d9d6b>", line 1, in <module>
    psi[:][:500] = b

ValueError: could not broadcast input array from shape (3,500) into shape (3,1000)

我对主题也有一些变化，但结果相似。这看起来非常简单。知道怎么做吗？

Answer 1

您可以使用：

idx = np.arange(3)
psi[idx, :500] = b[idx]

或者，如果循环中的

n

与

b

的第一个维度匹配：

psi[:, :500] = b

比较两种方法：

psi1 = np.zeros((3, 1000))
b = np.arange(3*500).reshape((3,500))

for n in range(3):
   psi1[n][:500] = b[n]

psi2 = np.zeros((3, 1000))
idx = np.arange(3)
psi2[idx, :500] = b[idx]

np.allclose(psi1, psi2)
# True