python之子集 2d numpy 数组并保持行一致

我想知道执行以下操作的最简单方法是什么:

假设我们有以下二维数组:

>>> a = np.array([['z', 'z', 'z', 'f', 'z','f', 'f'], ['z', 'z', 'z', 'f', 'z','f', 'f']]) 
 
array([['z', 'z', 'z', 'f', 'z', 'f', 'f'], 
   ['z', 'z', 'z', 'f', 'z', 'f', 'f']], 
  dtype='<U1') 
 
 
 
>>> b = np.array(range(0,14)).reshape(2, -1) 
 
 
array([[ 0,  1,  2,  3,  4,  5,  6], 
       [ 7,  8,  9, 10, 11, 12, 13]]) 
 
 
>>> idxs = list(zip(*np.where(a == 'f'))) 
 
[(0, 3), (0, 5), (0, 6), (1, 3), (1, 5), (1, 6)] 
 
 
>>> [b[x] for x in idxs] 
 
[3, 5, 6, 10, 12, 13] 

但是，我想保留之前关于第一个索引或行的结构 - 即:

[[3, 5, 6], [7, 11]] 

有没有办法轻松保持这种结构？

请您参考如下方法:

这是一个更复杂但纯 NumPy 的解决方案:

1. 获取索引(在 a 的扁平化版本中)是 'f'。
2. 获取新行开始的索引
3. 从1开始查找数组中属于一行的索引
4. 在这些索引处拆分数组。

代码看起来像这样:

>>> indices = np.flatnonzero(a.ravel() == 'f') 
>>> rows = np.arange(1, a.shape[0])*a.shape[1] 
>>> np.split(b.ravel()[indices], np.searchsorted(indices, rows)) 
[array([3, 5, 6], dtype=int64), array([10, 12, 13], dtype=int64)] 

比其他解决方案长一点，我不确定它是否会更快¹。

尽管就个人而言，我会选择列表理解和 zip:

[b_row[a_row] for a_row, b_row in zip(a == 'f', b)] 

它要短得多，而且根据我的时间安排，它的性能非常好。

---

时间:

import numpy as np 
a = np.array([['z', 'z', 'z', 'f', 'z','f', 'f']]*10000) 
b = np.arange(a.size).reshape(-1, a.shape[1]) 
 
%%timeit 
 
indices = np.flatnonzero(a.ravel() == 'f') 
rows = np.arange(1, a.shape[0])*a.shape[1] 
np.split(b.ravel()[indices], np.searchsorted(indices, rows)) 

123 ms ± 8.25 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit [b[i][a[i] == 'f'] for i in range(len(a))] 

162 ms ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

但与 my suggestion at Psidoms answer 相比要慢很多:

%timeit [b_row[a_row] for a_row, b_row in zip(a == 'f', b)] 

44.9 ms ± 1.93 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)