如何使用高级索引创建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])]   

- Merlin

---

使用broadcasting：np.array([ 1, 2, 3, 4, 16, 31]) + np.arange(10)[:,None]。 - Divakar

我需要一个变量y和aa - 这只是一个玩具例子。在现实世界中，我实际上正在切片“aa”。 - Merlin

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]])

- Divakar

---

我需要填充“aa”吗？如果偏移量的最后一个元素大于“aa”的长度，会发生错误还是会循环回来。 - Merlin

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×6的数据为[1,2,3,4,16,31]的矩阵。

如果您想让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]])

- Willem Van Onsem