在 numpy 数组中赋值

Question

我有一个 numpy 零数组。具体来说，假设它是 2x3x4：

x = np.zeros((2,3,4))

假设我有一个 2x3 的随机整数数组，从 0 到 3（x 的第 3 维的索引）。

>>> y = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))
>>> y
[[2 1 0]
 [3 2 0]]

我如何有效地执行以下分配（编辑：不使用 for 循环并且适用于 x 具有任意数量的维度和每个维度中任意数量的元素的东西）？

>>> x[0,0,y[0,0]]=1
>>> x[0,1,y[0,1]]=1
>>> x[0,2,y[0,2]]=1
>>> x[1,0,y[1,0]]=1
>>> x[1,1,y[1,1]]=1
>>> x[1,2,y[1,2]]=1
>>> x
array([[[ 0.,  0.,  1.,  0.],
        [ 0.,  1.,  0.,  0.],
        [ 1.,  0.,  0.,  0.]],

       [[ 0.,  0.,  0.,  1.],
        [ 0.,  0.,  1.,  0.],
        [ 1.,  0.,  0.,  0.]]])

谢谢， 詹姆斯

Answer 1

目前，我只能想到“简单”版本，其中涉及沿前两个维度进行展平。这段代码应该可以工作：

shape_last = x.shape[-1]
x.reshape((-1, shape_last))[np.arange(y.size), y.flatten()] = 1

这会产生（我随机生成的y）：

array([[[ 0.,  0.,  0.,  1.],
        [ 0.,  0.,  1.,  0.],
        [ 0.,  1.,  0.,  0.]],

       [[ 0.,  1.,  0.,  0.],
        [ 0.,  0.,  0.,  1.],
        [ 0.,  1.,  0.,  0.]]])

关键是，如果您使用多个 numpy 数组 (advanced indexing) 进行索引，numpy 将使用索引对来索引数组。

当然，请确保 x 和 y 都是 C 订单或 F 订单 - 否则，对 reshape 和 flatten 的调用可能会给出不同的订单。

Answer 2

使用numpy.meshgrid() 创建索引数组，您可以使用它们来索引原始数组和第三维值数组。

import numpy as np
import scipy as sp
import scipy.stats.distributions

a = np.zeros((2,3,4))
z = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))

xx, yy = np.meshgrid( np.arange(2), np.arange(3) )
a[ xx, yy, z[xx, yy] ] = 1
print a

为了清楚起见，我已将您的数组从 x 重命名为 a，并将索引数组从 y 重命名为 z。

编辑： 4D 示例：

a = np.zeros((2,3,4,5))
z = sp.stats.distributions.randint.rvs(0, 4, size=(2,3))
w = sp.stats.distributions.randint.rvs(0, 5, size=(2,3))

xx, yy = np.meshgrid( np.arange(2), np.arange(3) )
a[ xx, yy, z[xx, yy], w[xx, yy] ] = 1

Answer 3

x = np.zeros((2,3,4))
y=np.array([[2, 1, 0],[3, 2, 0]]) # or y=sp.stats...
for i in range(2):
    for j in range(3):
        x[i,j,y[i,j]]=1

将产生预期的结果，IIRC。如果数组维度永远不会改变，请考虑将两个 for 循环及其负担替换为

for j in range(3):
    x[0,j,y[0,j]] = x[1,j,y[1,j]] = 1