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

【问题标题】：Select multiple rows multiple times from a numpy array at once avoiding loops一次从 numpy 数组中多次选择多行，避免循环
【发布时间】：2021-02-12 09:51:25
【问题描述】：

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

鉴于M 和indexes，我希望N 避免for 循环，因为它对于大尺寸来说很慢。

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


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

【问题讨论】：

你得到一个形状不同的数组列表这一事实强烈表明这个列表理解是你能做的最好的。
Numpy 在处理同质数据时通常处于最佳状态，而您的预期输出却不是。循环似乎是这里的最佳选择。
@hpaul 的列表理解真的比np.split 更好吗？

标签： python arrays numpy numpy-ndarray numpy-slicing

【解决方案1】：

我们可以利用输出数据在至少一个维度上是同质的这一优势。

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

示例运行：

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

【讨论】：

对于这个小例子，简单的列表理解更快。但替代方案的规模可能不同。 split 也必须循环，取多个切片。缩放可能取决于行数与列数。

【解决方案2】：

使用广播的稍微不同的方法，以及识别分割点的不同方法：

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, 0],
       [1, 0, 0, 1, 1]]),
array([[1, 1, 1, 1, 0],
       [0, 0, 0, 1, 1],
       [1, 0, 0, 1, 1]]),
array([[0, 0, 0, 1, 1]]),
array([[1, 0, 1, 1, 0],
       [1, 1, 1, 1, 0],
       [1, 0, 0, 1, 1]])]

print (result_split)

【讨论】：