子集2d numpy數組並保持行一致

Question

我想知道執行以下操作的最簡單方法是：

假設我們有以下二維數組：

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

有沒有一種方法可以輕松保持此結構？

Answer 1

使用for循環：

[b[i][a[i] == 'f'] for i in range(len(a))]
# [array([3, 5, 6]), array([10, 12, 13])]

Answer 2

這是一個更復雜但純粹的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毫秒±8.25毫秒（平均±標准偏差，共7次運行，每個循環10次）

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

每個循環162 ms±14 ms（平均±標准偏差，共7次運行，每個循環1次）

但是，與我在Psidoms回答中的建議相比，速度要慢得多：

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

每個循環44.9 ms±1.93 ms（平均±標准偏差，共運行7次，每個循環10個）

Answer 3

 a = np.array([['z', 'z', 'z', 'f', 'z','f', 'f'], ['z', 'z', 'z', 'f', 'z','f', 'f']]) b = np.array(range(0,14)).reshape(2, -1) idxs = list(zip(*np.where(a == 'f'))) c=[[],[]] for x in idxs: c[x[0]].append(b[x]) print c 

Answer 4

In [89]: idx = np.where(a == 'f')
In [90]: idx
Out[90]: 
(array([0, 0, 0, 1, 1, 1], dtype=int32),
 array([3, 5, 6, 3, 5, 6], dtype=int32))

我們可以應用where元組在b選擇項目：

In [93]: b[idx]
Out[93]: array([ 3,  5,  6, 10, 12, 13])

等效地應用布爾掩碼：

In [94]: b[a == 'f']
Out[94]: array([ 3,  5,  6, 10, 12, 13])

np.argwhere進行where的轉置，產生一個2d數組，如idxs 。

In [95]: np.argwhere(a == 'f')
Out[95]: 
array([[0, 3],
       [0, 5],
       [0, 6],
       [1, 3],
       [1, 5],
       [1, 6]], dtype=int32)

如刪除與布爾型遮罩相對應的數組中的所有元素中所述 ，通常，我們不能選擇帶有遮罩的元素並保留某種2d結構。 在某些情況下，我們可以將1d結果重塑為有意義的東西。

In [96]: b[idx].reshape(2,-1)
Out[96]: 
array([[ 3,  5,  6],
       [10, 12, 13]])

逐行收集這些值並允許每一行具有不同大小結果的一種簡單方法是迭代：

In [100]: [j[i=='f'] for i,j in zip(a,b)]
Out[100]: [array([3, 5, 6]), array([10, 12, 13])]
In [101]: [j[i=='f'].tolist() for i,j in zip(a,b)]
Out[101]: [[3, 5, 6], [10, 12, 13]]