Numpy: récupère l'index des éléments d'un tableau 1d comme tableau 2d

J'ai un tableau numpy comme celui-ci: [1 2 2 0 0 1 3 5]

Est-il possible d'obtenir l'index des éléments sous forme de tableau 2D? Par exemple, la réponse pour l'entrée ci-dessus serait[[3 4], [0 5], [1 2], [6], [], [7]]

Actuellement, je dois boucler les différentes valeurs et appeler numpy.where(input == i)chaque valeur, qui a des performances terribles avec une entrée assez grande.

Voici une approche O (max (x) + len (x)) utilisant 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]]

Cela fonctionne en créant une matrice clairsemée avec des entrées aux positions (x [0], 0), (x [1], 1), ... CSC format (colonne creuse compressée) c'est assez simple. La matrice est ensuite convertie au LILformat (liste liée). Ce format stocke les indices de colonne pour chaque ligne sous forme de liste dans son rowsattribut, donc tout ce que nous devons faire est de prendre cela et de le convertir en liste.

Notez que pour les petites baies, les argsortsolutions basées sont probablement plus rapides, mais à une taille pas incroyablement grande, cela se croisera.

ÉDITER:

argsortbasée sur une numpyseule solution:

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 l'ordre des indices au sein des groupes n'a pas d'importance, vous pouvez également essayer argpartition(cela ne fait aucune différence dans ce petit exemple mais ce n'est pas garanti en général):

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

ÉDITER:

@Divakar déconseille l'utilisation de np.split. Au lieu de cela, une boucle est probablement plus rapide:

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 vous pouvez utiliser le tout nouvel opérateur de morse (Python3.8 +):

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

MODIFIER (MODIFIÉ):

(Pas pur numpy): Comme alternative à numba (voir le post de @ senderle), nous pouvons également utiliser pythran.

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

Ici numbagagne par une moustache en termes de performances:

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]

Trucs plus anciens:

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

Timings vs numba (ancien)

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

Une option potentielle en fonction de la taille de vos données consiste à simplement abandonner numpyet utiliser 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

Ensuite, vous vous retrouvez avec un dictionnaire de {value1: [index1, index2, ...], value2: [index3, index4, ...]}. L'échelle de temps est assez proche de la linéarité avec la taille du tableau, donc 10 000 000 prennent environ 2,7 secondes sur ma machine, ce qui semble assez raisonnable.

Bien que la demande concerne une numpysolution, j'ai décidé de voir s'il existe une numbasolution intéressante . Et en effet il y en a! Voici une approche qui représente la liste partitionnée sous la forme d'un tableau déchiqueté stocké dans un seul tampon préalloué. Cela s'inspire de l' argsortapproche proposée par Paul Panzer . (Pour une version plus ancienne qui ne fonctionnait pas aussi bien, mais qui était plus simple, voir ci-dessous.)

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

Cela traite une liste de dix millions d'éléments en 75 ms, ce qui représente une accélération de près de 50 fois par rapport à une version basée sur une liste écrite en pur Python.

Pour une version plus lente mais un peu plus lisible, voici ce que j'avais avant, basé sur un support expérimental récemment ajouté pour les "listes typées" de taille dynamique, qui nous permettent de remplir chaque bac de manière désordonnée beaucoup plus rapidement.

Cela lutte un peu avec numbale moteur d'inférence de type, et je suis sûr qu'il existe une meilleure façon de gérer cette partie. Cela s'avère également être presque 10 fois plus lent que ce qui précède.

@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

Je les ai testés par rapport aux éléments suivants:

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

Je les ai également testés contre une version cython précompilée similaire à enum_bins_numba_buffer(décrite en détail ci-dessous).

Sur une liste de dix millions d'ints aléatoires ( ints = np.random.randint(0, 100, 10000000)) j'obtiens les résultats suivants:

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)

Impressionnant, cette façon de travailler avec numbasurpasse une cythonversion de la même fonction, même avec la vérification des limites désactivée. Je n'ai pas encore assez de connaissances pythranpour tester cette approche en l'utilisant, mais je serais intéressé de voir une comparaison. Il semble probable sur la base de cette accélération que lepythran version pourrait également être un peu plus rapide avec cette approche.

Voici la cythonversion pour référence, avec quelques instructions de construction. Une fois que vous avez cythoninstallé, vous aurez besoin d'un setup.pyfichier simple comme celui-ci:

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

Et le module 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)]

Avec ces deux fichiers dans votre répertoire de travail, exécutez cette commande:

python setup.py build_ext --inplace

Vous pouvez ensuite importer la fonction à l'aide de from enum_bins_cython import enum_bins_cython.

Voici une façon vraiment vraiment bizarre de faire cela, c'est terrible, mais je l'ai trouvé trop drôle pour ne pas le partager - et tout 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: c'est la meilleure méthode que j'ai pu trouver le long de ce chemin. C'est toujours 10 fois plus lent que la solution 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)

Vous pouvez le faire en créant un dictionnaire de nombres, les clés seraient les nombres et les valeurs devraient être les indices que le nombre a vus, c'est l'un des moyens les plus rapides de le faire, vous pouvez voir le code ci-dessous:

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

Pseudocode:

1. obtenir le "nombre de tableaux 1d dans le tableau 2d", en soustrayant la valeur minimale de votre tableau numpy de la valeur maximale, puis plus un. Dans votre cas, ce sera 5-0 + 1 = 6

2. initialiser un tableau 2D avec le nombre de tableaux 1D qu'il contient. Dans votre cas, initialisez un tableau 2D avec 6 tableaux 1D. Chaque tableau 1d correspond à un élément unique de votre tableau numpy, par exemple, le premier tableau 1d correspondra à '0', le second tableau 1d correspondra à '1', ...

3. boucle dans votre tableau numpy, placez l'index de l'élément dans le tableau 1d correspondant droit. Dans votre cas, l'index du premier élément de votre tableau numpy sera placé dans le deuxième tableau 1d, l'indice du deuxième élément de votre tableau numpy sera mis dans le troisième tableau 1d, ....

Ce pseudocode prendra un temps linéaire pour s'exécuter car il dépend de la longueur de votre tableau numpy.

Cela vous donne exactement ce que vous voulez et prendrait environ 2,5 secondes pour 10 000 000 sur ma machine:

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

Donc, étant donné une liste d'éléments, vous voulez faire des paires (élément, index). En temps linéaire, cela pourrait se faire comme:

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

Cela devrait prendre du temps O (n). Je ne peux pas penser à une solution plus rapide pour l'instant, mais je mettrai à jour ici si je le fais.