在numpy數組中分配值

Question

我有一個numpy數組的零。 具體來說，假設它是2x3x4：

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

假設我有一個從0到3的2x3隨機整數數組（x的第三維的索引）。

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

如何有效地執行以下分配（編輯：不用於循環的東西，適用於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數組（ 高級索引 ）進行索引 ，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