Сравнение двух массивов 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]] 

Любые советы по этому поводу?

EDIT: Так что в конце концов это работает. Не очень удобно, но он работает.

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-массивы. Одно из предостережений заключается в том, что вам нужна глубокая копия первого массива a, так как np.view не смешивается с кусочками, ну. Но если это не слишком большая проблема для вас, что-то вдоль линий:

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.