¿Hay una función NumPy para devolver el primer índice de algo en una matriz?

Sé que hay un método para que una lista de Python devuelva el primer índice de algo:

>>> l = [1, 2, 3]
>>> l.index(2)
1

¿Hay algo así para las matrices NumPy?

Sí, aquí está la respuesta dada una matriz NumPy array, y un valor item, para buscar:

itemindex = numpy.where(array==item)

El resultado es una tupla con primero todos los índices de fila, luego todos los índices de columna.

Por ejemplo, si una matriz tiene dos dimensiones y contiene su artículo en dos ubicaciones, entonces

array[itemindex[0][0]][itemindex[1][0]]

sería igual a su artículo y también lo haría

array[itemindex[0][1]][itemindex[1][1]]

numpy.where

Si necesita el índice de la primera aparición de un solo valor , puede usar nonzero(o where, lo que equivale a lo mismo en este caso):

>>> t = array([1, 1, 1, 2, 2, 3, 8, 3, 8, 8])
>>> nonzero(t == 8)
(array([6, 8, 9]),)
>>> nonzero(t == 8)[0][0]
6

Si necesita el primer índice de cada uno de los muchos valores , obviamente podría hacer lo mismo que anteriormente varias veces, pero hay un truco que puede ser más rápido. A continuación se encuentran los índices del primer elemento de cada subsecuencia :

>>> nonzero(r_[1, diff(t)[:-1]])
(array([0, 3, 5, 6, 7, 8]),)

Observe que encuentra el comienzo de ambas subsecuencias de 3s y ambas subsecuencias de 8s:

[ 1 , 1, 1, 2 , 2, 3 , 8 , 3 , 8 , 8]

Por lo tanto, es ligeramente diferente a encontrar la primera aparición de cada valor. En su programa, puede trabajar con una versión ordenada tpara obtener lo que desea:

>>> st = sorted(t)
>>> nonzero(r_[1, diff(st)[:-1]])
(array([0, 3, 5, 7]),)

También puede convertir una matriz NumPy para listar en el aire y obtener su índice. Por ejemplo,

l = [1,2,3,4,5] # Python list
a = numpy.array(l) # NumPy array
i = a.tolist().index(2) # i will return index of 2
print i

Imprimirá 1.

Solo para agregar un muy eficiente y práctico numbaalternativa basada en np.ndenumerateencontrar el primer índice:

from numba import njit
import numpy as np

@njit
def index(array, item):
    for idx, val in np.ndenumerate(array):
        if val == item:
            return idx
    # If no item was found return None, other return types might be a problem due to
    # numbas type inference.

Esto es bastante rápido y trata naturalmente con matrices multidimensionales :

>>> arr1 = np.ones((100, 100, 100))
>>> arr1[2, 2, 2] = 2

>>> index(arr1, 2)
(2, 2, 2)

>>> arr2 = np.ones(20)
>>> arr2[5] = 2

>>> index(arr2, 2)
(5,)

Esto puede ser mucho más rápido (porque está cortocircuitando la operación) que cualquier enfoque que use np.whereo np.nonzero.

Sin embargo np.argwhere, también podría tratar con gracia las matrices multidimensionales (necesitaría convertirlo manualmente en una tupla y no está en cortocircuito), pero fallaría si no se encuentra ninguna coincidencia:

>>> tuple(np.argwhere(arr1 == 2)[0])
(2, 2, 2)
>>> tuple(np.argwhere(arr2 == 2)[0])
(5,)

Si va a usar esto como un índice en otra cosa, puede usar índices booleanos si las matrices son de escala amplia; No necesitas índices explícitos. La forma más simple de hacer esto es simplemente indexar en función de un valor de verdad.

other_array[first_array == item]

Cualquier operación booleana funciona:

a = numpy.arange(100)
other_array[first_array > 50]

El método distinto de cero también toma booleanos:

index = numpy.nonzero(first_array == item)[0][0]

Los dos ceros son para la tupla de índices (suponiendo que first_array es 1D) y luego el primer elemento en la matriz de índices.

l.index(x)devuelve el i más pequeño de manera que i es el índice de la primera aparición de x en la lista.

Se puede suponer con seguridad que la index()función en Python se implementa para que se detenga después de encontrar la primera coincidencia, y esto da como resultado un rendimiento promedio óptimo.

Para encontrar un elemento que se detiene después de la primera coincidencia en una matriz NumPy, use un iterador ( ndenumerate ).

In [67]: l=range(100)

In [68]: l.index(2)
Out[68]: 2

Matriz NumPy:

In [69]: a = np.arange(100)

In [70]: next((idx for idx, val in np.ndenumerate(a) if val==2))
Out[70]: (2L,)

Tenga en cuenta que ambos métodos index()y nextdevuelven un error si no se encuentra el elemento. Con next, se puede usar un segundo argumento para devolver un valor especial en caso de que no se encuentre el elemento, p. Ej.

In [77]: next((idx for idx, val in np.ndenumerate(a) if val==400),None)

Hay otras funciones en NumPy ( argmax, wherey nonzero) que puede ser utilizado para encontrar un elemento en una matriz, pero todos tienen el inconveniente de ir a través de todo el conjunto en busca de todas las ocurrencias, por tanto, no está optimizado para encontrar el primer elemento. Tenga en cuenta también eso wherey nonzerodevuelva las matrices, por lo que debe seleccionar el primer elemento para obtener el índice.

In [71]: np.argmax(a==2)
Out[71]: 2

In [72]: np.where(a==2)
Out[72]: (array([2], dtype=int64),)

In [73]: np.nonzero(a==2)
Out[73]: (array([2], dtype=int64),)

Comparación de tiempo

Simplemente comprobando que para matrices grandes la solución que usa un iterador es más rápida cuando el elemento buscado está al comienzo de la matriz (usando %timeiten el shell de IPython):

In [285]: a = np.arange(100000)

In [286]: %timeit next((idx for idx, val in np.ndenumerate(a) if val==0))
100000 loops, best of 3: 17.6 µs per loop

In [287]: %timeit np.argmax(a==0)
1000 loops, best of 3: 254 µs per loop

In [288]: %timeit np.where(a==0)[0][0]
1000 loops, best of 3: 314 µs per loop

Este es un problema abierto de NumPy GitHub .

Ver también: Numpy: encuentre el primer índice de valor rápidamente

Para las matrices ordenadas unidimensionales , sería mucho más simple y eficiente O (log (n)) usar numpy.searchsorted que devuelve un entero NumPy (posición). Por ejemplo,

arr = np.array([1, 1, 1, 2, 3, 3, 4])
i = np.searchsorted(arr, 3)

Solo asegúrate de que la matriz ya esté ordenada

También verifique si el índice devuelto realmente contiene el elemento buscado, ya que el objetivo principal de searchsorted es encontrar índices donde se deben insertar elementos para mantener el orden.

if arr[i] == 3:
    print("present")
else:
    print("not present")

Para indexar cualquier criterio, puede hacer algo como lo siguiente:

In [1]: from numpy import *
In [2]: x = arange(125).reshape((5,5,5))
In [3]: y = indices(x.shape)
In [4]: locs = y[:,x >= 120] # put whatever you want in place of x >= 120
In [5]: pts = hsplit(locs, len(locs[0]))
In [6]: for pt in pts:
   .....:         print(', '.join(str(p[0]) for p in pt))
4, 4, 0
4, 4, 1
4, 4, 2
4, 4, 3
4, 4, 4

Y aquí hay una función rápida para hacer lo que hace list.index (), excepto que no genera una excepción si no se encuentra. Cuidado, esto es probablemente muy lento en matrices grandes. Probablemente pueda mono parchear esto en matrices si prefiere usarlo como método.

def ndindex(ndarray, item):
    if len(ndarray.shape) == 1:
        try:
            return [ndarray.tolist().index(item)]
        except:
            pass
    else:
        for i, subarray in enumerate(ndarray):
            try:
                return [i] + ndindex(subarray, item)
            except:
                pass

In [1]: ndindex(x, 103)
Out[1]: [4, 0, 3]

Para los arreglos 1D, recomendaría np.flatnonzero(array == value)[0], lo que es equivalente a ambos np.nonzero(array == value)[0][0]y np.where(array == value)[0][0]evita la fealdad de desempaquetar una tupla de 1 elemento.

Una alternativa para seleccionar el primer elemento de np.where () es usar una expresión generadora junto con enumerate, como:

>>> import numpy as np
>>> x = np.arange(100)   # x = array([0, 1, 2, 3, ... 99])
>>> next(i for i, x_i in enumerate(x) if x_i == 2)
2

Para una matriz bidimensional, uno haría:

>>> x = np.arange(100).reshape(10,10)   # x = array([[0, 1, 2,... 9], [10,..19],])
>>> next((i,j) for i, x_i in enumerate(x) 
...            for j, x_ij in enumerate(x_i) if x_ij == 2)
(0, 2)

La ventaja de este enfoque es que deja de verificar los elementos de la matriz después de encontrar la primera coincidencia, mientras que np.where verifica todos los elementos para encontrar una coincidencia. Una expresión generadora sería más rápida si hay una coincidencia temprana en la matriz.

Hay muchas operaciones en NumPy que tal vez podrían reunirse para lograr esto. Esto devolverá índices de elementos iguales al elemento:

numpy.nonzero(array - item)

Luego, puede tomar los primeros elementos de las listas para obtener un solo elemento.

El paquete numpy_indexed (descargo de responsabilidad, soy su autor) contiene un equivalente vectorizado de list.index para numpy.ndarray; es decir:

sequence_of_arrays = [[0, 1], [1, 2], [-5, 0]]
arrays_to_query = [[-5, 0], [1, 0]]

import numpy_indexed as npi
idx = npi.indices(sequence_of_arrays, arrays_to_query, missing=-1)
print(idx)   # [2, -1]

Esta solución ha vectorizado el rendimiento, se generaliza a ndarrays y tiene varias formas de tratar con valores perdidos.

Nota: esto es para la versión Python 2.7

Puede usar una función lambda para tratar el problema, y funciona tanto en la matriz NumPy como en la lista.

your_list = [11, 22, 23, 44, 55]
result = filter(lambda x:your_list[x]>30, range(len(your_list)))
#result: [3, 4]

import numpy as np
your_numpy_array = np.array([11, 22, 23, 44, 55])
result = filter(lambda x:your_numpy_array [x]>30, range(len(your_list)))
#result: [3, 4]

Y puedes usar

result[0]

para obtener el primer índice de los elementos filtrados.

Para python 3.6, use

list(result)

en vez de

result