如何提取第一個元素位於不同列表的第二個數組中的那些行？

Question

假設我有這個數組

array([[100, 1], 
     [200, 2], 
     [300, 3], 
     [400, 4], 
     [440, 3]]) 

而且我有這個列表或一維數組[100,300]。 我想要我的操作輸出[1,3]。 如何在numpy中做到這一點。

我實際上是在Theano中使用這些numpy數組（一個使用gpu加速計算的機器學習庫）。我將有很多行。 Numpy數組允許我無縫地將它們用作Theano中的Tensor對象。但是如果我不得不使用字典，那麼我必須用普通的Python來做到這一點，而且我不確定一旦我轉向大數據時它是否能夠保持良好。所以我實際上是在尋找一個numpy操作，索引中的一些技巧或類似的東西。

Answer 1

你可以使用np.in1d：如果您需要致電np.in1d爲vals許多不同的值

In [12]: arr 
Out[12]: 
array([[100, 1], 
     [200, 2], 
     [300, 3], 
     [400, 4], 
     [440, 3]]) 

In [14]: vals = [100, 300] 
In [23]: np.in1d(arr[:,0], vals) 
Out[23]: array([ True, False, True, False, False], dtype=bool) 

In [24]: arr[np.in1d(arr[:,0], vals), 1] 
Out[24]: array([1, 3]) 

，那麼它可以支付準備一個字典作爲arshajii建議，因爲在準備字典後（一個O(n)操作，其中n = len(arr)），擡頭看e值將是O(m)操作，其中m = len(vals)。

如果n變得非常大，但是字典可能需要太多內存。在這種情況下，您可能需要使用np.in1d。

如果指數（鍵）值都是ints和小幅度的，還有你可以用它來獲取O(m)性能，而無需使用字典一NumPy的索引招：

In [30]: big = np.full(arr[:,0].max()+1, np.nan) 

In [31]: big[arr[:,0]] = arr[:,1] 

In [32]: big[vals] 
Out[32]: array([ 1., 3.]) 

準備big是O(n)操作，但索引big[vals]是O(m)。如果arr[:,0].max()很小，並且關鍵值是ints，則使用big的優點是所需內存少於使用dict的內存。

---

In [33]: %timeit arr[np.in1d(arr[:,0], vals), 1] 
10000 loops, best of 3: 21.5 µs per loop 

In [34]: %timeit big[vals] 
1000000 loops, best of 3: 1.23 µs per loop 

與arshajii的解決方案比較：

In [38]: d = dict(arr) 
In [40]: %timeit [d[k] for k in vals] 
1000000 loops, best of 3: 447 ns per loop 

所以用最好的方法取決於arr和vals，大小你多少次被執行此操作，多少內存你有，如果鑰匙很小ints。您需要對與您的用例相關的數據進行基準測試，以做出正確的決定。

Answer 2

我想你的陣列簡單地轉換成詞典：

>>> a = array([[100, 1], 
...   [200, 2], 
...   [300, 3], 
...   [400, 4], 
...   [440, 3]]) 
>>> 
>>> keys = [100, 300] 
>>> 
>>> d = dict(a) 
>>> 
>>> [d[k] for k in keys] 
[1, 3] 

Answer 3

如果您確定要搜索的所有值實際上都存在於搜索陣列中，您還可以使用np.searchsorted。對於大型陣列，與其他建議相比似乎更快。

s = np.sort(A[:,0]) 
A[np.searchsorted(s, values), 1] 

如果要搜索的數組已經排序，您可以省略排序關閉過程，操作更加快捷。