2d 數組作為 3d 數組的索引

Question

我有一個 8000x64 元素的二維數組 (C)、一個 8000x1 元素的一維數組 (s) 和另一個 1x64 元素的一維數組 (d)。 索引 i 的每一行，其中 s[i] 為 True，應由向量 d 添加。 這很有效：

C[s == True] += d

現在我已經為 C、s 和 d 添加了一個維度，上面的邏輯將應用於附加維度的每個元素。

下面的代碼做了我想要的，但速度很慢。

for i in range(I):
        C_this = C[:,:,i]
        s_this = s[:,i]
        d_this = d[:,i]

        C_this[s_this == True] += d_this
        C[:,:,i] = C_this

有沒有一種沒有 for 循環的 numpy 方法來做到這一點？

Answer 1

開始時使用額外維度會更容易：

In [376]: C = np.zeros((4,2,3),int)                                                            
In [377]: s = np.array([[0,0],[0,1],[1,0],[1,1]],bool)                                         
In [378]: d = np.arange(1,13).reshape(4,3)                                                     
In [379]: C.shape, s.shape, d.shape                                                            
Out[379]: ((4, 2, 3), (4, 2), (4, 3))
In [380]: I,J = np.nonzero(s)                                                                  
In [381]: I,J                                                                                  
Out[381]: (array([1, 2, 3, 3]), array([1, 0, 0, 1]))

In [383]: C[I,J]=d[I]                                                                          
In [384]: C                                                                                    
Out[384]: 
array([[[ 0,  0,  0],
        [ 0,  0,  0]],

       [[ 0,  0,  0],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 0,  0,  0]],

       [[10, 11, 12],
        [10, 11, 12]]])

你的方式：

In [385]: C = np.zeros((4,2,3),int)                                                            
In [386]: for i in range(4): 
     ...:     C[i,:,:][s[i,:]] += d[i,:] 
     ...:                                                                                      
In [387]: C                                                                                    
Out[387]: 
array([[[ 0,  0,  0],
        [ 0,  0,  0]],

       [[ 0,  0,  0],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 0,  0,  0]],

       [[10, 11, 12],
        [10, 11, 12]]])

Answer 2

由於 numpy 索引的工作方式， s在第一個示例中選擇C的相關行。 要在 3D 情況下做同樣的事情，您必須將C重塑為(8000*3, 64)並將s重塑為(8000*3, 1) 。 現在唯一的問題是讓d考慮每個第三維中不同的行數，這可以通過np.repeat完成。

第一部分是

C2 = np.swapaxes(C, -1, 1).reshape(-1, 64)

這是非常低效的，因為它會復制整個數組。 更好的安排是，如果C以(3, 8000, 64)開頭。 然后您只需要解開前兩個軸即可獲得正確的形狀和內存布局，而無需復制數據。

repeats = np.count_nonzero(s, axis=0)
C.reshape(-1, 64)[s.ravel()] += np.repeat(d, repeats, axis=0)

由於 reshape 操作在這種情況下返回一個視圖，索引應該正常工作以就地遞增。 不過，我認為這種方法不一定非常好，因為它復制d每一行的次數與s在新維度的相應元素中不為零的次數相同。

Answer 3

這是我對@hpaulj 提出的方法的實現。 請注意，我不想從他那里獲得功勞，所以請將他的答案標記為正確，而不是我的答案。 只是想分享我所做的。

import numpy as np
import numpy.random as npr

C = np.zeros((100, 8000, 64), dtype=int)
s = np.zeros((100, 8000), dtype=bool)
d = np.zeros((100, 64), dtype=int)

C[:,:,:] = npr.randint(50, size=C.shape)
s[:,:] = npr.randint(3, size=s.shape)
d[:,:] = npr.randint(10, size=d.shape)

I, J = np.nonzero(s)
C[I, J] += d[I]

然后我分析了我制作的程序，它在我的機器上運行的時間不到 450 毫秒（最后兩行不到 300 毫秒）。 請注意，對“randint”的調用只是為了設置數組值，因此這些行不適用於您的用例。