Numpy: obtém o índice dos elementos de uma matriz 1d como uma matriz 2d

Eu tenho uma matriz numpy como esta: [1 2 2 0 0 1 3 5]

É possível obter o índice dos elementos como uma matriz 2D? Por exemplo, a resposta para a entrada acima seria[[3 4], [0 5], [1 2], [6], [], [7]]

Atualmente, tenho que repetir os diferentes valores e chamar numpy.where(input == i)cada valor, que tem um desempenho terrível com uma entrada grande o suficiente.

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 rowsatributo; portanto, tudo o que precisamos fazer é pegar isso e convertê-lo em lista.

Observe que, para pequenas matrizes, as argsortsoluções baseadas são provavelmente mais rápidas, mas em um tamanho não insanamente grande, isso passará.

EDITAR:

argsortnumpyú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 numbavence 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

Uma opção potencial, dependendo do tamanho dos seus dados, é desistir numpye 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

Então você acaba com um dicionário de {value1: [index1, index2, ...], value2: [index3, index4, ...]}. A escala de tempo é quase linear com o tamanho da matriz, então 10.000.000 demoram ~ 2.7s na minha máquina, o que parece bastante razoável.

Embora o pedido seja uma numpysolução, decidi ver se há uma numbasolução interessante . E de fato existe! Aqui está uma abordagem que representa a lista particionada como uma matriz irregular armazenada em um único buffer pré-alocado. Isso se inspira na argsortabordagem proposta por Paul Panzer . (Para uma versão mais antiga que não foi tão bem, mas era mais simples, veja abaixo.)

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

Isso processa uma lista de dez milhões de itens em 75ms, que é quase uma aceleração de 50x de uma versão baseada em lista escrita em Python puro.

Para uma versão mais lenta, mas um pouco mais legível, eis o que eu tinha antes, com base no suporte experimental recentemente adicionado a "listas digitadas" de tamanho dinâmico, que nos permitem encher cada lixeira de maneira fora de ordem muito mais rapidamente.

Isso luta numbaum pouco com o mecanismo de inferência de tipos e tenho certeza de que há uma maneira melhor de lidar com essa parte. Isso também é quase 10 vezes mais lento que o acima.

@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

Eu testei estes contra o seguinte:

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

Também os testei em uma versão cython pré-compilada semelhante a enum_bins_numba_buffer(descrita em detalhes abaixo).

Em uma lista de dez milhões de ints aleatórios ( ints = np.random.randint(0, 100, 10000000)), obtenho os seguintes 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)

De maneira impressionante, essa maneira de trabalhar numbasupera uma cythonversão da mesma função, mesmo com a verificação de limites desativada. Ainda não tenho familiaridade suficiente pythranpara testar essa abordagem, mas estaria interessado em ver uma comparação. Parece provável, com base nessa aceleração, que a pythranversão também seja um pouco mais rápida com essa abordagem.

Aqui está a cythonversão para referência, com algumas instruções de construção. Depois de cythoninstalar, você precisará de um setup.pyarquivo simples 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()]
)

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

Com esses dois arquivos em seu diretório de trabalho, execute este comando:

python setup.py build_ext --inplace

Você pode importar a função usando from enum_bins_cython import enum_bins_cython.

Aqui está uma maneira realmente muito estranha de fazer isso que é terrível, mas achei engraçado demais para não compartilhar - e tudo 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: este é o melhor método que eu poderia encontrar nesse caminho. Ainda é 10 vezes mais lento que a argsortsolução da @PaulPanzer :

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)

Você pode fazer isso criando um dicionário de números, as chaves seriam os números e os valores deveriam ser os índices que o número visualizado, esta é uma das maneiras mais rápidas de fazer isso, você pode ver o código abaixo:

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

Pseudo-código:

1. obtenha o "número de matrizes 1d na matriz 2d", subtraindo o valor mínimo da sua matriz numpy do valor máximo e depois mais um. No seu caso, será 5-0 + 1 = 6

2. inicialize uma matriz 2d com o número de matrizes 1d dentro dela. No seu caso, inicialize uma matriz 2d com 6 1d. Cada matriz 1d corresponde a um elemento único na sua matriz numpy, por exemplo, a primeira matriz 1d corresponde a '0', a segunda matriz 1d corresponde a '1', ...

3. loop através de sua matriz numpy, coloque o índice do elemento na matriz 1d correspondente à direita. No seu caso, o índice do primeiro elemento na sua matriz numpy será colocado na segunda matriz 1d, o índice do segundo elemento na sua matriz numpy será colocado na terceira matriz 1d, ....

Esse pseudocódigo levará um tempo linear para ser executado, pois depende do comprimento da sua matriz numpy.

Isso fornece exatamente o que você deseja e levaria cerca de 2,5 segundos para 10.000.000 na minha 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()))

Portanto, dada uma lista de elementos, você deseja criar pares (elemento, índice). Em tempo linear, isso pode ser feito 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())

Isso deve levar tempo O (n). Não consigo pensar em uma solução mais rápida a partir de agora, mas atualizarei aqui se o fizer.