2015-10-26 2 views
4

Если у меня есть Numpy массив так:Numpy: поиск первой совпадающей строки

import numpy as np 
x = np.array([[0,1],[0,2],[1,1],[0,2]]) 

Как я могу вернуть индекс первой строки, которая соответствует [0,2]?

Для списков это легко с помощью index:

[[0,1],[0,2],[1,1],[0,2]] 

l.index([0,2]) 
> 1 

И я знаю, что numpy имеет функцию numpy.where, но я не совсем уверен, что сделать из выхода numpy.where:

np.where(x==[0,2]) 
> (array([0, 1, 1, 3, 3]), array([0, 0, 1, 0, 1])) 

Там также numpy.argmax, но это также не возвращает то, что я ищу, что является просто индексом 1

np.argmax(x == [0,2], axis = 1) 

ответ

4

Если список поиска [0,2], Вы доведете в broadcasting по сравнению с x, что дает нам маску той же формы, как x. Поскольку вы ищете точное соответствие, вы будете искать строки со всеми значениями TRUE с .all(1). Наконец, вам нужен первый индекс, поэтому используйте np.where или np.nonzero и выберите первый элемент. Реализация в качестве пробных прогонов будет:

In [132]: x 
Out[132]: 
array([[0, 1], 
     [0, 2], 
     [1, 1], 
     [0, 2]]) 

In [133]: search_list = [0,2] 

In [134]: np.where((x == search_list).all(1))[0][0] 
Out[134]: 1 

In [135]: np.nonzero((x == search_list).all(1))[0][0] 
Out[135]: 1 
Смежные вопросы