一次从numpy数组中多次选择多行，避免循环

如何解决一次从numpy数组中多次选择多行，避免循环

我正在寻找一种在给定索引数组的情况下从numpy数组中多次选择多行的方法。

考虑到M和indexes，我想让N避免for循环，因为在大尺寸情况下它很慢。

import numpy as np
M = np.array([[1,1,0],[1,[0,1],1]])
indexes = np.array([[True,False,True],[False,True,False],[True,True]])
N = [M[index] for index in indexes]


Out: 
[array([[1,1]]),array([[1,array([[0,1]])]

解决方法

我们可以利用输出数据在至少一维上均匀的优势。

x,y = np.where(indexes)
split_idx = np.flatnonzero(np.diff(x))+1
output = np.split(M[y],split_idx)

样品运行：

>>> x
array([0,1,2,3,3],dtype=int32)
>>> y
array([0,dtype=int32)
>>> split_idx
array([2,5,6],dtype=int32)

,

使用广播的方法略有不同，并且使用不同的方法来识别分裂点：

b_shape = (indexes.shape[0],) + M.shape  # New shape for broadcasted M. Here,(4,4,5)
M_b = np.broadcast_to(M,b_shape)        # Broadcasted M with the new shape.
                                         # (it uses views instead of replicating data)
r,c = np.nonzero(indexes)
result_joined = M_b[r,c,:]                             # The stack of all the selected rows from M
split_points = np.cumsum(np.sum(indexes,axis=1))[:-1] # Identify where to split.
result_split = np.split (result_merged,split_points)  # Final result,obtained by splitting.

输出：

[array([[1,0],[1,1]]),array([[1,[0,1],array([[0,1]])]

print (result_split)