Aqui está uma abordagem 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]]
Isso funciona criando uma matriz esparsa com entradas nas posições (x [0], 0), (x [1], 1), ... Usando o formato CSC
(coluna esparsa compactada) isso é bastante simples. A matriz é então convertida para o formato LIL
(lista vinculada). Esse formato armazena os índices da coluna para cada linha como uma lista em seu rows
atributo; portanto, tudo o que precisamos fazer é pegar isso e convertê-lo em lista.
Observe que, para pequenas matrizes, as argsort
soluções baseadas são provavelmente mais rápidas, mas em um tamanho não insanamente grande, isso passará.
EDITAR:
argsort
numpy
única solução baseada em :
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 a ordem dos índices dentro dos grupos não importar, você também pode tentar argpartition
(isso não faz diferença neste pequeno exemplo, mas isso não é garantido em geral):
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 recomenda contra o uso de np.split
. Em vez disso, um loop é provavelmente mais 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))]
Ou você pode usar o novo 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):
(Não é puro numpy): Como alternativa ao numba (consulte a publicação de @ senderle), também podemos usar o pythran.
Ajuntar com 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)]
Aqui numba
vence por um bigode em termos de desempenho:
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]
Coisas mais antigas:
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
Horários vs. numba (antigo)
timeit(lambda:bincollect(x),number=10)
# 3.5732191529823467
timeit(lambda:enumerate_bins(x),number=10)
# 6.7462647299980745
np.argsort([1, 2, 2, 0, 0, 1, 3, 5])
dáarray([3, 4, 0, 5, 1, 2, 6, 7], dtype=int64)
. então você pode apenas comparar os próximos elementos.