Numpy: ottiene l'indice degli elementi di un array 1d come un array 2d

Ho un array intorpidito come questo: [1 2 2 0 0 1 3 5]

È possibile ottenere l'indice degli elementi come un array 2d? Ad esempio, la risposta per l'input sopra sarebbe[[3 4], [0 5], [1 2], [6], [], [7]]

Attualmente devo eseguire il loop dei diversi valori e chiamare numpy.where(input == i)per ciascun valore, che ha prestazioni terribili con un input abbastanza grande.

Ecco un approccio 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]]

Funziona creando una matrice sparsa con voci nelle posizioni (x [0], 0), (x [1], 1), ... Usando il formato CSC(colonna sparsa compressa) questo è piuttosto semplice. La matrice viene quindi convertita nel formato LIL(elenco collegato). Questo formato memorizza gli indici di colonna per ogni riga come elenco nel suo rowsattributo, quindi tutto ciò che dobbiamo fare è prenderlo e convertirlo in elenco.

Si noti che per le matrici di piccole dimensioni le argsortsoluzioni sono probabilmente più veloci, ma in alcuni casi di dimensioni non follemente grandi, questo andrà oltre.

MODIFICARE:

argsortbasata su numpysoluzione -Solo:

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])]

Se l'ordine degli indici all'interno dei gruppi non ha importanza, puoi anche provare argpartition(succede che non fa alcuna differenza in questo piccolo esempio, ma questo non è garantito in generale):

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])]

MODIFICARE:

@Divakar sconsiglia l'uso di np.split. Invece, un ciclo è probabilmente più veloce:

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

Oppure potresti usare il nuovissimo operatore tricheco (Python3.8 +):

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

EDIT (Montaggio):

(Non puro intorpidimento): in alternativa a numba (vedi il post di @ senderle) possiamo anche usare pythran.

Compila 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)]

Qui numbavince da un baffo dal punto di vista delle prestazioni:

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]

Roba precedente:

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

Tempi vs. numba (vecchio)

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

Un'opzione potenziale a seconda della dimensione dei dati è semplicemente abbandonare numpye utilizzare 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

Quindi si finisce con un dizionario di {value1: [index1, index2, ...], value2: [index3, index4, ...]}. Il ridimensionamento del tempo è abbastanza vicino al lineare con le dimensioni dell'array, quindi 10.000.000 impiegano circa 2,7 secondi sulla mia macchina, il che sembra abbastanza ragionevole.

Sebbene la richiesta sia per una numpysoluzione, ho deciso di vedere se esiste una numbasoluzione interessante . E davvero c'è! Ecco un approccio che rappresenta l'elenco partizionato come un array irregolare memorizzato in un singolo buffer preallocato. Questo si ispira argsortall'approccio proposto da Paul Panzer . (Per una versione precedente che non funzionava altrettanto bene, ma era più semplice, vedi sotto).

@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)]

Questo elabora un elenco di dieci milioni di articoli in 75ms, che è quasi un 50 volte più veloce rispetto a una versione basata su elenco scritta in puro Python.

Per una versione più lenta ma un po 'più leggibile, ecco quello che avevo prima, basato sul supporto sperimentale recentemente aggiunto per "elenchi tipizzati" di dimensioni dinamiche, che ci consentono di riempire ogni cestino in modo fuori servizio molto più rapidamente.

Questo lotta numbaun po ' con il motore dell'inferenza del tipo, e sono sicuro che c'è un modo migliore per gestire quella parte. Anche questo risulta essere quasi 10 volte più lento di quanto sopra.

@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

Ho testato questi contro quanto segue:

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()

Li ho anche testati con una versione cython precompilata simile a enum_bins_numba_buffer(descritta in dettaglio di seguito).

Su un elenco di dieci milioni di ints casuali ( ints = np.random.randint(0, 100, 10000000)) ottengo i seguenti risultati:

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)

Incredibilmente, questo modo di lavorare numbasupera una cythonversione della stessa funzione, anche con il controllo dei limiti disattivato. Non ho ancora abbastanza familiarità pythranper testare questo approccio utilizzandolo, ma sarei interessato a vedere un confronto. Sembra probabilmente basato su questo speedup che la pythranversione potrebbe anche essere un po 'più veloce con questo approccio.

Ecco la cythonversione di riferimento, con alcune istruzioni di compilazione. Una volta cythoninstallato, avrai bisogno di un semplice setup.pyfile come questo:

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()]
)

E il modulo 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 questi due file nella directory di lavoro, esegui questo comando:

python setup.py build_ext --inplace

È quindi possibile importare la funzione utilizzando from enum_bins_cython import enum_bins_cython.

Ecco un modo davvero strano per farlo, è terribile, ma l'ho trovato troppo divertente per non condividere - e tutto 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]]

EDIT: questo è il metodo migliore che ho potuto trovare lungo questo percorso. È ancora 10 volte più lento della soluzione di @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)

Puoi farlo creando un dizionario di numeri, le chiavi sarebbero i numeri e i valori dovrebbero essere gli indici di quel numero visto, questo è uno dei modi più veloci per farlo, puoi vedere il codice qui sotto:

>>> 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] 

pseudocodice:

1. ottenere il "numero di array 1d nell'array 2d" sottraendo il valore minimo dell'array numpy dal valore massimo e quindi più uno. Nel tuo caso, sarà 5-0 + 1 = 6

2. inizializza un array 2d con il numero di array 1d al suo interno. Nel tuo caso, inizializza un array 2d con 6 array 1d al suo interno. Ogni array 1d corrisponde a un elemento univoco nell'array numpy, ad esempio, il primo array 1d corrisponderà a '0', il secondo array 1d corrisponderà a '1', ...

3. scorrere attraverso l'array numpy, inserire l'indice dell'elemento nell'array 1d corrispondente a destra. Nel tuo caso, l'indice del primo elemento nell'array numpy verrà inserito nel secondo array 1d, l'indice del secondo elemento nell'array numpy verrà inserito nel terzo array 1d, ....

L'esecuzione di questo pseudocodice richiederà tempo lineare poiché dipende dalla lunghezza dell'array numpy.

Questo ti dà esattamente quello che vuoi e richiederebbe circa 2,5 secondi per 10.000.000 sulla mia macchina:

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()))

Quindi, dato un elenco di elementi, si desidera creare coppie (elemento, indice). In tempo lineare, questo potrebbe essere fatto come:

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())

Questo dovrebbe richiedere O (n) tempo. Al momento non riesco a pensare a una soluzione più veloce, ma aggiornerò qui se lo faccio.