比較兩個numpy數組並添加相同的行

Question

我有兩個大的數據文件，一個有兩列，一個有三列。我想從第二個文件中包含在第一個數組中的所有行中選擇。我的想法是比較numpy陣列。

比方說，我有：

a = np.array([[1, 2, 3], [3, 4, 5], [1, 4, 6]]) 

b = np.array([[1, 2], [3, 4]]) 

，結果應該是這樣的：

[[1, 2, 3], [3, 4, 5]] 

上有什麼建議？編輯： 所以最終這個作品。不太方便，但它的作品。

for ii in range(a.shape[0]): 
    u, v, w = a[ii,:] 
    for jj in range(b.shape[0]): 
     if (u == b[jj, 0] and v == b[jj, 1]): 
      print [u, v, w] 

Answer 1

請檢查，@Ernie。它可以幫助您找到解決方案。 ; d

http://docs.scipy.org/doc/numpy-1.10.1/reference/generated/numpy.in1d.html

Answer 2

NumPy的數組不是絕對必要在這裏，你可以用列表來代替：

>>> a = [[1, 2, 3], [3, 4, 5], [1, 4, 6]] 
>>> b = [[1, 2], [3, 4]] 
>>> [row for row in a if row[:2] in b] 
[[1, 2, 3], [3, 4, 5]] 

上面的解決方案可以耗費時間的大名單。如果您有效率問題，我認爲@Eelco答案將是首選。但是，如果你不想安裝numpy_indexed（或根本不能），你會發現這個代碼有用：

>>> a = np.array([[1, 2, 3], [3, 4, 5], [1, 4, 6]]) 
>>> b = np.array([[1, 2], [3, 4]]) 
>>> b_set = set(map(tuple, b)) 
>>> b_set 
set([(1, 2), (3, 4)]) 
>>> a[np.nonzero([tuple(row[:-1]) in b_set for row in a])] 
array([[1, 2, 3], 
     [3, 4, 5]]) 

Answer 3

的numpy_indexed包（免責聲明：我是它的作者）包含的功能來解決這樣的問題高效，不使用任何蟒蛇循環：

import numpy_indexed as npi 
a[npi.contains(b, a[:, :2])] 

Answer 4

如果您不想使用另一個庫，但要做到這一點的只有numpy，你可以做類似的東西有什麼建議here和here，即使用np.in1d（見docs） e用掩碼錶示一個1D數組中的元素是否存在於另一個1D數組中。正如名稱所示，此功能僅適用於1D陣列。但是你可以使用結構化數組視圖（使用np.view）來欺騙numpy，以爲你有1D數組。但有一點需要注意，因爲np.view不會與切片混合，所以您需要第一個陣列的深層副本a。但是，如果這對你來說不是太大的問題，可以採取以下行動：

a_cp = a[:, :2].copy() 
a[np.in1d(a_cp.view((np.void, a_cp.dtype.itemsize*a_cp.shape[1])).ravel(), 
      b.view((np.void, b.dtype.itemsize*b.shape[1])).ravel())] 

可能適合你。

這直接使用掩碼陣列從您的數組a返回正確的值。