python - 我可以用掩码分割 numpy 数组吗？

Question

我想将数组拆分为带有掩码和索引的数组
像下面

a = array([ 0,  1,  2,  3,  4, 5]))  
b = [0,2,3]  

进入

c = array([[0, 2, 3], [1, 3, 4], [2, 4, 5]])  

我可以在没有循环的情况下执行此操作吗？

编辑:

更多例子...

假设我们有一个数组 a带形状 [10, 10, 10]哪里 a[x, y, :] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
现在给出面具 b = [0, 3, 7]
我希望输出是一个数组 c带形状 [10, 10, 3, 3]哪里 c[x, y, :, :] = [[0, 3, 7], [1, 4, 8], [2, 5, 9]]

Answer 1

您可以生成 b作为您想要的索引和移位向量之间的广播总和。然后你可以再次广播成更大的尺寸。由于示例中的输出不依赖于 a数组，我不理会。

from numpy import array, broadcast_to, arange
from numpy.random import random

a = random((10,10,10)) # not used on the code at all.... don't understand what it is for...

b = [0,2,3]
b_array = array(b)
b_shifts = arange(3).reshape(-1,1)
c_cell= b+b_shifts # here, they are broadcasted toegether. one is a row-vector and one is a column-vector...
c = broadcast_to(c_cell,(10,10,3,3))

您可能想要创建 b_shifts使用其他方法取决于步长等......

编辑
根据您的评论，似乎更准确的答案是:

from numpy import array, arange
a = arange(2*2*10).reshape((2,2,10)) # some example input 
b = array([0,2,3])                   # the 'template' to extract
shifts = arange(3).reshape(-1,1)     # 3 is the number of repeats
indexer = b+shifts                   # broadcasted sum makes a matrix
c = a[:,:,indexer]                   # extract

这将需要 b数组作为一种模板，并以一定的移位重复它。最后，它将从每个数组 a[i,j,:] 中提取这些条目。进入 c[i,j,:,:] .上面的输出是:
print(a)

[[[ 0  1  2  3  4  5  6  7  8  9]
  [10 11 12 13 14 15 16 17 18 19]]
 [[20 21 22 23 24 25 26 27 28 29]
  [30 31 32 33 34 35 36 37 38 39]]]

print(c)

[[[[ 0  2  3]
   [ 1  3  4]
   [ 2  4  5]]
  [[10 12 13]
   [11 13 14]
   [12 14 15]]]
 [[[20 22 23]
   [21 23 24]
   [22 24 25]]
  [[30 32 33]
   [31 33 34]
   [32 34 35]]]]