我可以用掩码拆分 numpy 数组吗？答案

【问题标题】：can I split numpy array with mask?我可以用掩码拆分 numpy 数组吗？
【发布时间】：2020-07-07 22:17:57
【问题描述】：

我想将数组拆分为带有掩码和索引的数组
如下所示

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

【问题讨论】：

我不明白您的示例输入/输出。你能想出另一个更清楚的例子吗？
这只是一个小例子，但我想在大小约为 100*300*300 的数组上执行此操作。循环会花费很多时间。有没有其他方法可以做到这一点？
如果a=array([9,4,7,2,4,3,9,7,5]) 和b=[1,1,3] 会发生什么？
您能否添加更多有关该问题的一般信息？ @LudvigH OP 想要访问a 中每个位置的第一项、第三项和第四项（由b 中的索引给出）。要回答您的问题，它会提供[[9, 9, 7], [4, 4, 2], [7, 7, 4], ...] 供您输入。

标签： python python-3.x numpy numpy-ndarray numpy-slicing

【解决方案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...

编辑根据您的 cmets，似乎更准确的答案是：

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

【讨论】：

谢谢，输入数组更大，c_cells 需要在 b 上循环。但如果没有别的办法，我就用这个。
@JYWon 更新了答案以考虑更长的班次。
谢谢！也许c = a[:, :, c_cell]是我想要的答案
顺便说一句，从您对 stride 的评论看来，您这样做是卷积中的一个步骤。您可以考虑直接使用 scipy.signal.fftconvolve 之类的方法来提高性能并减少样板代码

【解决方案2】：

使用as_strided的移动窗口方法：

In [1]: a = np.arange(6)                                                                       
In [2]: a                                                                                      
Out[2]: array([0, 1, 2, 3, 4, 5])
In [3]: as_strided = np.lib.stride_tricks.as_strided

对于这个单班，strides 参数很简单。 shape 需要更多思考——我们期望有多少行，以及最大索引：

In [5]: b = as_strided(a, shape=(3,4), strides=(8,8))                                          
In [6]: b                                                                                      
Out[6]: 
array([[0, 1, 2, 3],
       [1, 2, 3, 4],
       [2, 3, 4, 5]])

然后选择列：

In [8]: b[:,[0,2,3]]                                                                           
Out[8]: 
array([[0, 2, 3],
       [1, 3, 4],
       [2, 4, 5]])

要将其扩展到 3d 案例，我将从 https://stackoverflow.com/a/60881930/901925 工作，LudvigH 的回答

In [10]: a = np.arange(2*2*10).reshape((2,2,10)) # some example input  
    ...: b = np.array([0,2,3])                                                                 
In [11]: a                                                                                     
Out[11]: 
array([[[ 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]]])
In [12]: a.shape                                                                               
Out[12]: (2, 2, 10)
In [13]: a.strides                                                                             
Out[13]: (160, 80, 8)
In [18]: a1 = as_strided(a, shape=(2,2,3,4), strides=(160,80,8,8))                             
In [19]: a1                                                                                    
Out[19]: 
array([[[[ 0,  1,  2,  3],
         [ 1,  2,  3,  4],
         [ 2,  3,  4,  5]],

        [[10, 11, 12, 13],
         [11, 12, 13, 14],
         [12, 13, 14, 15]]],


       [[[20, 21, 22, 23],
         [21, 22, 23, 24],
         [22, 23, 24, 25]],

        [[30, 31, 32, 33],
         [31, 32, 33, 34],
         [32, 33, 34, 35]]]])

这只是第一个案例的扩展，前 2 个维度只是随行。这是展开到 2d 窗口中的最后一个维度。

再次选择列的子集：

In [20]: a1[:,:,:,b]                                                                           
Out[20]: 
array([[[[ 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]]]])

as_strided 步骤很有效，创建了一个view。但是索引会复制。我们必须花一些时间来针对 LudvigH 的移位索引方法对其进行测试。

【讨论】：