與NumPy的較短的一個襯裏用於抓取索引

Question

我有下面的代碼行：

idxs = [i for i,x in enumerate(labels) if x==lbl] 

• 標籤是整數的numpy的陣列
• LBL是int

idxs =索引st相應的標籤元素有價值lbl

問題：是否有一個較短的單線？

謝謝！

Answer 1

你可以使用的numpy.where的一個參數形式：

idxs = np.where(labels == lbl)[0] 

，或者等價地，使用numpy.nonzero：（！感謝喬）

idxs = np.nonzero(labels == lbl)[0] 

，或者更好的可讀性，

idxs = np.flatnonzero(labels == lbl) 

---

例如，

In [332]: np.random.seed(1) 

In [333]: labels = np.random.randint(5, size=10) 

In [334]: labels 
Out[334]: array([3, 4, 0, 1, 3, 0, 0, 1, 4, 4]) 

In [335]: [i for i,x in enumerate(labels) if x==lbl] 
Out[335]: [3, 7] 

In [336]: np.where(labels == lbl)[0] 
Out[336]: array([3, 7]) 

---

使用np.where是多少，比列表修真大型陣列更快：

In [339]: labels = np.tile(labels, 1000) 

In [340]: labels.shape 
Out[340]: (10000,) 

In [341]: %timeit np.where(labels == lbl)[0] 
10000 loops, best of 3: 45.9 µs per loop 

In [342]: %timeit [i for i,x in enumerate(labels) if x==lbl] 
100 loops, best of 3: 5.31 ms per loop 

In [343]: 5310/45.9 
Out[343]: 115.68627450980392 

Answer 2

我沒有代表。以評論答案...然而，請記住，當使用numpy.where「標籤」必須是一個numpy數組。

Codelifting unutbu的回答是：

idxs = np.where(np.array(labels) == lbl)[0]

只是要清楚：正確的答案被unutbu製造。