Numpy: obtenga el índice de los elementos de una matriz 1d como una matriz 2d

Tengo una matriz numpy como esta: [1 2 2 0 0 1 3 5]

¿Es posible obtener el índice de los elementos como una matriz 2d? Por ejemplo, la respuesta para la entrada anterior sería[[3 4], [0 5], [1 2], [6], [], [7]]

Actualmente tengo que recorrer los diferentes valores y solicitar numpy.where(input == i)cada valor, que tiene un rendimiento terrible con una entrada lo suficientemente grande.

Aquí hay un enfoque O (max (x) + len (x)) usando scipy.sparse:

import numpy as np
from scipy import sparse

x = np.array("1 2 2 0 0 1 3 5".split(),int)
x
# array([1, 2, 2, 0, 0, 1, 3, 5])


M,N = x.max()+1,x.size
sparse.csc_matrix((x,x,np.arange(N+1)),(M,N)).tolil().rows.tolist()
# [[3, 4], [0, 5], [1, 2], [6], [], [7]]

Esto funciona creando una matriz dispersa con entradas en las posiciones (x [0], 0), (x [1], 1), ... Usando el CSCformato (columna dispersa comprimida) esto es bastante simple. La matriz se convierte luego al LILformato (lista vinculada). Este formato almacena los índices de columna para cada fila como una lista en su rowsatributo, por lo que todo lo que tenemos que hacer es tomar eso y convertirlo a la lista.

Tenga en cuenta que las argsortsoluciones basadas en matrices pequeñas son probablemente más rápidas, pero en algunos de tamaño no increíblemente grande, esto se cruzará.

EDITAR:

argsortbasada solo en numpysolución:

np.split(x.argsort(kind="stable"),np.bincount(x)[:-1].cumsum())
# [array([3, 4]), array([0, 5]), array([1, 2]), array([6]), array([], dtype=int64), array([7])]

Si el orden de los índices dentro de los grupos no importa, también puede intentarlo argpartition(no hace ninguna diferencia en este pequeño ejemplo, pero esto no está garantizado en general):

bb = np.bincount(x)[:-1].cumsum()
np.split(x.argpartition(bb),bb)
# [array([3, 4]), array([0, 5]), array([1, 2]), array([6]), array([], dtype=int64), array([7])]

EDITAR:

@Divakar recomienda contra el uso de np.split. En cambio, un bucle es probablemente más rápido:

A = x.argsort(kind="stable")
B = np.bincount(x+1).cumsum()
[A[B[i-1]:B[i]] for i in range(1,len(B))]

O puede utilizar el nuevo operador de morsa (Python3.8 +):

A = x.argsort(kind="stable")
B = np.bincount(x)
L = 0
[A[L:(L:=L+b)] for b in B.tolist()]

EDITAR (EDITADO):

(No puro numpy): como alternativa a numba (ver la publicación de @ senderle) también podemos usar pythran.

Compilar con pythran -O3 <filename.py>

import numpy as np

#pythran export sort_to_bins(int[:],int)

def sort_to_bins(idx, mx):
    if mx==-1: 
        mx = idx.max() + 1
    cnts = np.zeros(mx + 2, int)
    for i in range(idx.size):
        cnts[idx[i] + 2] += 1
    for i in range(3, cnts.size):
        cnts[i] += cnts[i-1]
    res = np.empty_like(idx)
    for i in range(idx.size):
        res[cnts[idx[i]+1]] = i
        cnts[idx[i]+1] += 1
    return [res[cnts[i]:cnts[i+1]] for i in range(mx)]

Aquí numbagana por un bigote en cuanto al rendimiento:

repeat(lambda:enum_bins_numba_buffer(x),number=10)
# [0.6235917090671137, 0.6071486569708213, 0.6096088469494134]
repeat(lambda:sort_to_bins(x,-1),number=10)
# [0.6235359431011602, 0.6264424560358748, 0.6217901279451326]

Cosas más antiguas:

import numpy as np

#pythran export bincollect(int[:])

def bincollect(a):
    o = [[] for _ in range(a.max()+1)]
    for i,j in enumerate(a):
        o[j].append(i)
    return o

Tiempos contra numba (antiguo)

timeit(lambda:bincollect(x),number=10)
# 3.5732191529823467
timeit(lambda:enumerate_bins(x),number=10)
# 6.7462647299980745

Una opción potencial dependiendo del tamaño de sus datos es simplemente abandonar numpyy usar collections.defaultdict:

In [248]: from collections import defaultdict

In [249]: d = defaultdict(list)

In [250]: l = np.random.randint(0, 100, 100000)

In [251]: %%timeit
     ...: for k, v in enumerate(l):
     ...:     d[v].append(k)
     ...:
10 loops, best of 3: 22.8 ms per loop

Entonces terminas con un diccionario de {value1: [index1, index2, ...], value2: [index3, index4, ...]}. La escala de tiempo es bastante lineal con el tamaño de la matriz, por lo que 10,000,000 toma ~ 2.7s en mi máquina, lo que parece bastante razonable.

Aunque la solicitud es para una numpysolución, decidí ver si hay una numbasolución basada en algo interesante . Y de hecho lo hay! Aquí hay un enfoque que representa la lista particionada como una matriz irregular almacenada en un único búfer preasignado. Esto se inspira en el argsortenfoque propuesto por Paul Panzer . (Para una versión anterior que no funcionó tan bien, pero que era más simple, ver más abajo).

@numba.jit(numba.void(numba.int64[:], 
                      numba.int64[:], 
                      numba.int64[:]), 
           nopython=True)
def enum_bins_numba_buffer_inner(ints, bins, starts):
    for x in range(len(ints)):
        i = ints[x]
        bins[starts[i]] = x
        starts[i] += 1

@numba.jit(nopython=False)  # Not 100% sure this does anything...
def enum_bins_numba_buffer(ints):
    ends = np.bincount(ints).cumsum()
    starts = np.empty(ends.shape, dtype=np.int64)
    starts[1:] = ends[:-1]
    starts[0] = 0

    bins = np.empty(ints.shape, dtype=np.int64)
    enum_bins_numba_buffer_inner(ints, bins, starts)

    starts[1:] = ends[:-1]
    starts[0] = 0
    return [bins[s:e] for s, e in zip(starts, ends)]

Esto procesa una lista de diez millones de elementos en 75 ms, que es casi un 50 veces más rápido que una versión basada en listas escrita en Python puro.

Para una versión más lenta pero algo más legible, esto es lo que tenía antes, basado en el soporte experimental recientemente agregado para "listas escritas" de tamaño dinámico, que nos permiten llenar cada contenedor de una manera desordenada mucho más rápidamente.

Esto lucha numbaun poco con el motor de inferencia de tipos, y estoy seguro de que hay una mejor manera de manejar esa parte. Esto también resulta ser casi 10 veces más lento que el anterior.

@numba.jit(nopython=True)
def enum_bins_numba(ints):
    bins = numba.typed.List()
    for i in range(ints.max() + 1):
        inner = numba.typed.List()
        inner.append(0)  # An awkward way of forcing type inference.
        inner.pop()
        bins.append(inner)

    for x, i in enumerate(ints):
        bins[i].append(x)

    return bins

Probé esto contra lo siguiente:

def enum_bins_dict(ints):
    enum_bins = defaultdict(list)
    for k, v in enumerate(ints):
        enum_bins[v].append(k)
    return enum_bins

def enum_bins_list(ints):
    enum_bins = [[] for i in range(ints.max() + 1)]
    for x, i in enumerate(ints):
        enum_bins[i].append(x)
    return enum_bins

def enum_bins_sparse(ints):
    M, N = ints.max() + 1, ints.size
    return sparse.csc_matrix((ints, ints, np.arange(N + 1)),
                             (M, N)).tolil().rows.tolist()

También los probé con una versión cython precompilada similar a enum_bins_numba_buffer(descrita en detalle a continuación).

En una lista de diez millones de entradas aleatorias ( ints = np.random.randint(0, 100, 10000000)) obtengo los siguientes resultados:

enum_bins_dict(ints)
3.71 s ± 80.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_list(ints)
3.28 s ± 52.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_sparse(ints)
1.02 s ± 34.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_numba(ints)
693 ms ± 5.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

enum_bins_cython(ints)
82.3 ms ± 1.77 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

enum_bins_numba_buffer(ints)
77.4 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Impresionantemente, esta forma de trabajar numbasupera a una cythonversión de la misma función, incluso con la comprobación de límites desactivada. Todavía no tengo suficiente familiaridad pythranpara probar este enfoque al usarlo, pero me interesaría ver una comparación. Parece probable en base a esta aceleración que elpythran versión también podría ser bastante más rápida con este enfoque.

Aquí está la cythonversión de referencia, con algunas instrucciones de compilación. Una vez que haya cythoninstalado, necesitará un setup.pyarchivo simple como este:

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy

ext_modules = [
    Extension(
        'enum_bins_cython',
        ['enum_bins_cython.pyx'],
    )
]

setup(
    ext_modules=cythonize(ext_modules),
    include_dirs=[numpy.get_include()]
)

Y el módulo Cython, enum_bins_cython.pyx:

# cython: language_level=3

import cython
import numpy
cimport numpy

@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef void enum_bins_inner(long[:] ints, long[:] bins, long[:] starts) nogil:
    cdef long i, x
    for x in range(len(ints)):
        i = ints[x]
        bins[starts[i]] = x
        starts[i] = starts[i] + 1

def enum_bins_cython(ints):
    assert (ints >= 0).all()
    # There might be a way to avoid storing two offset arrays and
    # save memory, but `enum_bins_inner` modifies the input, and
    # having separate lists of starts and ends is convenient for
    # the final partition stage.
    ends = numpy.bincount(ints).cumsum()
    starts = numpy.empty(ends.shape, dtype=numpy.int64)
    starts[1:] = ends[:-1]
    starts[0] = 0

    bins = numpy.empty(ints.shape, dtype=numpy.int64)
    enum_bins_inner(ints, bins, starts)

    starts[1:] = ends[:-1]
    starts[0] = 0
    return [bins[s:e] for s, e in zip(starts, ends)]

Con estos dos archivos en su directorio de trabajo, ejecute este comando:

python setup.py build_ext --inplace

Luego puede importar la función usando from enum_bins_cython import enum_bins_cython.

Aquí hay una manera realmente extraña de hacer esto que es terrible, pero me pareció demasiado divertido para no compartirlo, ¡y todo numpy!

out = np.array([''] * (x.max() + 1), dtype = object)
np.add.at(out, x, ["{} ".format(i) for i in range(x.size)])
[[int(i) for i in o.split()] for o in out]

Out[]:
[[3, 4], [0, 5], [1, 2], [6], [], [7]]

EDITAR: este es el mejor método que pude encontrar en este camino. Todavía es 10 veces más lento que la solución de @PaulPanzer argsort:

out = np.empty((x.max() + 1), dtype = object)
out[:] = [[]] * (x.max() + 1)
coords = np.empty(x.size, dtype = object)
coords[:] = [[i] for i in range(x.size)]
np.add.at(out, x, coords)
list(out)

Puede hacerlo haciendo un diccionario de números, las claves serían los números y los valores deberían ser los índices que ese número ve, esta es una de las formas más rápidas de hacerlo, puede ver el código a continuación:

>>> import numpy as np
>>> a = np.array([1 ,2 ,2 ,0 ,0 ,1 ,3, 5])
>>> b = {}
# Creating an empty list for the numbers that exist in array a
>>> for i in range(np.min(a),np.max(a)+1):
    b[str(i)] = []

# Adding indices to the corresponding key
>>> for i in range(len(a)):
    b[str(a[i])].append(i)

# Resulting Dictionary
>>> b
{'0': [3, 4], '1': [0, 5], '2': [1, 2], '3': [6], '4': [], '5': [7]}

# Printing the result in the way you wanted.
>>> for i in sorted (b.keys()) :
     print(b[i], end = " ")

[3, 4] [0, 5] [1, 2] [6] [] [7] 

Pseudocódigo:

1. obtenga el "número de matrices 1d en la matriz 2d", restando el valor mínimo de su matriz numpy del valor máximo y luego más uno. En su caso, será 5-0 + 1 = 6

2. Inicializar una matriz 2d con el número de matrices 1d dentro de ella. En su caso, inicialice una matriz 2d con 6 matrices 1d. Cada matriz 1d corresponde a un elemento único en su matriz numpy, por ejemplo, la primera matriz 1d corresponderá a '0', la segunda matriz 1d corresponderá a '1', ...

3. recorra su matriz numpy, coloque el índice del elemento en la matriz 1d correspondiente. En su caso, el índice del primer elemento en su matriz numpy se colocará en la segunda matriz 1d, el índice del segundo elemento en su matriz numpy se colocará en la tercera matriz 1d ...

Este pseudocódigo tardará un tiempo lineal en ejecutarse, ya que depende de la longitud de su matriz numpy.

Esto le da exactamente lo que desea y tomaría alrededor de 2.5 segundos por 10,000,000 en mi máquina:

import numpy as np
import timeit

# x = np.array("1 2 2 0 0 1 3 5".split(),int)
x = np.random.randint(0, 100, 100000)

def create_index_list(x):
    d = {}
    max_value = -1
    for i,v in enumerate(x):
        if v > max_value:
            max_value = v
        try:
            d[v].append(i)
        except:
            d[v] = [i]
    result_list = []
    for i in range(max_value+1):
        if i in d:
            result_list.append(d[i])
        else:
            result_list.append([])
    return result_list

# print(create_index_list(x))
print(timeit.timeit(stmt='create_index_list(x)', number=1, globals=globals()))

Entonces, dada una lista de elementos, desea hacer pares (elemento, índice). En tiempo lineal, esto podría hacerse como:

hashtable = dict()
for idx, val in enumerate(mylist):
    if val not in hashtable.keys():
         hashtable[val] = list()
    hashtable[val].append(idx)
newlist = sorted(hashtable.values())

Esto debería llevar tiempo O (n). No puedo pensar en una solución más rápida a partir de ahora, pero actualizaré aquí si lo hago.