如何使用花式索引創建Numpy數組

Question

如何使用numpy的花式索引來創建這個，我想要最快的性能：

array([[ 1,  2,  3,  4, 16, 31],
       [ 2,  3,  4,  5, 17, 32],
       [ 3,  4,  5,  6, 18, 33],
       [ 4,  5,  6,  7, 19, 34],
       [ 5,  6,  7,  8, 20, 35],
       [ 6,  7,  8,  9, 21, 36],
       [ 7,  8,  9, 10, 22, 37],
       [ 8,  9, 10, 11, 23, 38],
       [ 9, 10, 11, 12, 24, 39],
       [10, 11, 12, 13, 25, 40]]

從這開始：

   a   = np.arange(0,10)
   aa  = np.arange(0,50)
   y   = 1
   AA  = [(aa[np.array([x+y, 1+x+y, 2+x+y, 3+x+y,  15+x+y, 30+x+y])]) for x, i in enumerate(a)]

我明白了，

[array([ 2,  3,  4,  5, 17, 32]),
 array([ 3,  4,  5,  6, 18, 33]),
 array([ 4,  5,  6,  7, 19, 34]),
 array([ 5,  6,  7,  8, 20, 35]),
 array([ 6,  7,  8,  9, 21, 36]),
 array([ 7,  8,  9, 10, 22, 37]),
 array([ 8,  9, 10, 11, 23, 38]),
 array([ 9, 10, 11, 12, 24, 39]),
 array([10, 11, 12, 13, 25, 40]),
 array([11, 12, 13, 14, 26, 41])]   

Answer 1

利用給定的變量a ， aa和y ，這里使用broadcasting進行外部添加 -

offset = np.array([0,1,2,3,15,30])
out = aa[a[:,None] + offset + y]

使用add-ufunc的顯式外部方法 -

out = aa[np.add.outer(a , offset + y)]

---

出界情況

對於超出界限的情況（aa小於要求），我們可以用aa填充零，然后將其索引到其中 -

offset = np.array([0,1,2,3,15,30])
idx = np.add.outer(a , offset + y)
aa_p = np.pad(aa,(0,idx.max()-len(a)+1), 'constant')
out = aa_p[idx]

或初始化輸出數組，然后創建一個有效位置的掩碼進行分配 -

offset = np.array([0,1,2,3,15,30])
idx = np.add.outer(a , offset + y)
mask = idx < len(aa)
out = np.zeros(idx.shape, dtype=aa.dtype)
out[mask] = aa[idx[mask]]

樣本輸入，輸出 -

In [234]: a   = np.arange(0,10)
     ...: aa  = np.arange(4,42)
     ...: y   = 6
     ...: 

In [235]: out
Out[235]: 
array([[10, 11, 12, 13, 25, 40],
       [11, 12, 13, 14, 26, 41],
       [12, 13, 14, 15, 27,  0],
       [13, 14, 15, 16, 28,  0],
       [14, 15, 16, 17, 29,  0],
       [15, 16, 17, 18, 30,  0],
       [16, 17, 18, 19, 31,  0],
       [17, 18, 19, 20, 32,  0],
       [18, 19, 20, 21, 33,  0],
       [19, 20, 21, 22, 34,  0]])

Answer 2

我們可以在這里使用廣播：

>>> np.arange(0,10).reshape(-1,1) + np.array([*range(1,5),16,31])
array([[ 1,  2,  3,  4, 16, 31],
       [ 2,  3,  4,  5, 17, 32],
       [ 3,  4,  5,  6, 18, 33],
       [ 4,  5,  6,  7, 19, 34],
       [ 5,  6,  7,  8, 20, 35],
       [ 6,  7,  8,  9, 21, 36],
       [ 7,  8,  9, 10, 22, 37],
       [ 8,  9, 10, 11, 23, 38],
       [ 9, 10, 11, 12, 24, 39],
       [10, 11, 12, 13, 25, 40]])

在這里，我們創建一個10×1矩陣，范圍從0到（不包括）10，我們用數據[1,2,3,4,16,31]創建一個1×6矩陣。

如果你想讓y成為“偏移量”，你可以把它寫成：

>>> y = 1
>>> np.arange(y,y+10).reshape(-1,1) + np.array([*range(0,4),15,30])
array([[ 1,  2,  3,  4, 16, 31],
       [ 2,  3,  4,  5, 17, 32],
       [ 3,  4,  5,  6, 18, 33],
       [ 4,  5,  6,  7, 19, 34],
       [ 5,  6,  7,  8, 20, 35],
       [ 6,  7,  8,  9, 21, 36],
       [ 7,  8,  9, 10, 22, 37],
       [ 8,  9, 10, 11, 23, 38],
       [ 9, 10, 11, 12, 24, 39],
       [10, 11, 12, 13, 25, 40]])