Самый эффективный способ получить указанные строки из 2-го массива? - PullRequest
Новая
       13

Самый эффективный способ получить указанные строки из 2-го массива?

5 голосов
/ 31 марта 2011

У меня есть двумерный массив с более чем 100 000 строк.Мне нужно вернуть подмножество этих строк (и мне нужно выполнять эти операции много тысяч раз, поэтому эффективность важна).

Пример макета выглядит следующим образом: 

import numpy as np
a = np.array([[1,5.5],
             [2,4.5],
             [3,9.0],
             [4,8.01]])
b = np.array([2,4])

Итак ... Я хочу вернуть массив из a со строками, указанными в первом столбце с помощью b: 

c=[[2,4.5],
   [4,8.01]]

Разница, конечно, состоит в том, что в обоихи б, поэтому я хотел бы избежать зацикливания.Кроме того, я играл с созданием словаря и использованием np.nonzero, но все еще немного озадачен.

Заранее спасибо за любые идеи!

РЕДАКТИРОВАТЬ: Обратите внимание, что,в этом случае b являются идентификаторами, а не индексами.Вот пересмотренный пример: 

import numpy as np
a = np.array([[102,5.5],
             [204,4.5],
             [343,9.0],
             [40,8.01]])
b = np.array([102,343])

И я хочу вернуть: 

c = [[102,5.5],
     [343,9.0]]

Ответы [ 2 ]

6 голосов
/ 31 марта 2011

РЕДАКТИРОВАТЬ: Удален мой первоначальный ответ, так как это было неправильное понимание вопроса. Вместо этого попробуйте: 

ii = np.where((a[:,0] - b.reshape(-1,1)) == 0)[1]
c = a[ii,:]

Я использую широковещательную передачу для вычитания каждого элемента b из a, а затем ищу нули в этом массиве, которые указывают на совпадение. Это должно сработать, но вы должны быть немного осторожны при сравнении чисел с плавающей запятой, особенно если b не является массивом целых чисел.

РЕДАКТИРОВАТЬ 2 Благодаря предложению Свена, вы можете попробовать эту слегка измененную версию вместо: 

ii = np.where(a[:,0] == b.reshape(-1,1))[1]
c = a[ii,:]

Это немного быстрее, чем моя первоначальная реализация.

РЕДАКТИРОВАТЬ 3 Самое быстрое решение на сегодняшний день (~ в 10 раз быстрее, чем второе решение Свена для больших массивов): 

c = a[np.searchsorted(a[:,0],b),:]

Предполагается, что a[:,0] отсортировано и все значения b отображаются в a[:,0].

4 голосов
/ 01 апреля 2011

Немного более краткий способ сделать это 

c = a[(a[:,0] == b[:,None]).any(0)]

Применяются обычные предостережения для сравнений с плавающей запятой.

Редактировать : Если b не слишком мало, следующее слегка странное решение работает лучше: 

b.sort()
c = a[b[np.searchsorted(b, a[:, 0]) - len(b)] == a[:,0]]
...