Encuentra los índices de elementos no nulos y agrupa por valores.

Escribí un código en python que toma una matriz numpy como entrada y devuelve una lista de índices agrupados por los valores correspondientes (es decir, la salida [3] devuelve todos los índices con valor de 3). Sin embargo, carezco del conocimiento de escribir código vectorizado y tuve que hacerlo usando ndenumerate. Esta operación solo tomó unos 9 segundos, lo cual es demasiado lento.

La segunda idea que tuve fue usar numpy.nonzero de la siguiente manera:

for i in range(1, max_value): current_array = np.nonzero(input == i) # save in an array 

Esto tomó 5,5 segundos y, por lo tanto, fue una buena mejora, pero aún así lenta. ¿Alguna forma de hacerlo sin bucles o forma optimizada de obtener los pares de índices por valor?

Aquí hay un algoritmo O (n log n) para su problema. La solución de bucle obvia es O (n), por lo que para conjuntos de datos suficientemente grandes, esto será más lento:

 >>> a = np.random.randint(3, size=10) >>> a array([1, 2, 2, 0, 1, 0, 2, 2, 1, 1]) >>> index = np.arange(len(a)) >>> sort_idx = np.argsort(a) >>> cnt = np.bincount(a) >>> np.split(index[sort_idx], np.cumsum(cnt[:-1])) [array([3, 5]), array([0, 4, 8, 9]), array([1, 2, 6, 7])] 

Dependerá del tamaño de sus datos, pero es razonablemente rápido para grandes conjuntos de datos:

 In [1]: a = np.random.randint(1000, size=1e6) In [2]: %%timeit ...: indices = np.arange(len(a)) ...: sort_idx = np.argsort(a) ...: cnt = np.bincount(a) ...: np.split(indices[sort_idx], np.cumsum(cnt[:-1])) ...: 10 loops, best of 3: 140 ms per loop 

Si está dispuesto a usar algo de memoria adicional, puede vectorizar mediante transmisión:

 import numpy as np input = np.random.randint(1,max_value, 100) indices = np.arange(1, max_value) matches = input == indices[:,np.newaxis] # broadcasts across each index 

Luego, las coincidencias para cada índice i son simplemente np.nonzero(matches[i]) .