2015-08-24 20 views
9

Czy istnieje sposób na uzyskanie indeksów kilku elementów w tablicy NumPy jednocześnie?Uzyskiwanie indeksów kilku elementów w tablicy NumPy naraz

E.g.

import numpy as np 
a = np.array([1, 2, 4]) 
b = np.array([1, 2, 3, 10, 4]) 

chciałbym znaleźć indeks każdego elementu a w b, a mianowicie: [0,1,4].

znajdę rozwiązanie używam trochę rozwlekły:

import numpy as np 

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

c = np.zeros_like(a) 
for i, aa in np.ndenumerate(a): 
    c[i] = np.where(b==aa)[0] 

print('c: {0}'.format(c)) 

wyjściowa:

c: [0 1 4] 
+1

np. Gdzie (np.in1d ​​(b, a)) zwraca (tablica ([0, 1, 4], dtype = int64),). Na podstawie ostatniego przykładu na http://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html. –

Odpowiedz

12

można użyć in1d i nonzero (lub where dla tej sprawy):

>>> np.in1d(b, a).nonzero()[0] 
array([0, 1, 4]) 

Działa to dobrze dla twoich przykładowych tablic, ale ogólnie rzecz biorąc promień zwróconych indeksów nie uwzględnia kolejności wartości w a. Może to być problem w zależności od tego, co chcesz zrobić dalej.

W takim przypadku znacznie lepiej odpowiedź jest jedna @Jaime daje here korzystając searchsorted:

>>> sorter = np.argsort(b) 
>>> sorter[np.searchsorted(b, a, sorter=sorter)] 
array([0, 1, 4]) 

ta zwraca indeksy dla wartości, które pojawiają się w a. Na przykład:

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

>>> sorter = np.argsort(b) 
>>> sorter[np.searchsorted(b, a, sorter=sorter)] 
array([3, 1, 0]) # the other method would return [0, 1, 3] 
2

Jest to prosta-liner przy użyciu pakietu numpy-indexed (Zastrzeżenie: Jestem jej autor):

import numpy_indexed as npi 
idx = npi.indices(b, a) 

Realizacja jest w pełni wektorowy, a to daje kontrolę nad obsługa brakujących wartości. Co więcej, działa również na tablicach nd (na przykład, odnajdując indeksy wierszy a w b).

Powiązane problemy