一次從 numpy 數組中多次選擇多行避免循環

Question

我正在尋找一種方法，可以在給定索引數組的情況下多次從 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]])]

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

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