Python >> Tutoriel Python >  >> Python Tag >> NumPy

Le moyen le plus efficace d'ajouter des éléments compte tenu de la liste d'index dans numpy

Je doute que vous puissiez aller beaucoup plus vite que np.bincount - et notez comment la documentation officielle fournit ce cas d'utilisation exact

# Your example
A = [0.5, 0.6]
D = [[0.1, 0.1, 0.2], [0.2, 0.4, 0.1]]
I = [[0, 1, 0], [0, 1, 1]]

# Solution
import numpy as np    
D, I = np.array(D).flatten(), np.array(I).flatten()
print(np.bincount(I, D)) #[0.5 0.6]

La forme de I et D n'a pas d'importance :vous pouvez clairement parcourir les tableaux sans changer le résultat :

index = np.ravel(I)
data = np.ravel(D)

Vous pouvez maintenant trier les deux tableaux selon I :

sorter = np.argsort(index)
index = index[sorter]
data = data[sorter]

C'est utile car maintenant index ressemble à ceci :

0, 0, 0, 1, 1, 1

Et data est-ce :

0.1, 0.2, 0.2, 0.1, 0.4, 0.1

L'addition de séries de nombres consécutifs devrait être plus simple que le traitement d'emplacements aléatoires. Commençons par trouver les indices où les exécutions commencent :

runs = np.r_[0, np.flatnonzero(np.diff(index)) + 1]

Maintenant, vous pouvez utiliser le fait que les ufuncs comme np.add avoir un reduce partiel opération appelée reduceat . Cela vous permet de sommer les régions d'un tableau :

a = np.add.reduceat(data, runs)

Si I est garanti pour contenir tous les indices dans [0, A.size ) au moins une fois, vous avez terminé :attribuez simplement à A au lieu de a . Sinon, vous pouvez faire le mappage en utilisant le fait que le début de chaque exécution en index est l'indice cible :

A = np.zeros(n)
A[index[runs]] = a

Analyse de la complexité algorithmique :

  • ravel est O(1) dans le temps et dans l'espace si les données sont dans un tableau. Si c'est une liste, c'est O(MN) dans le temps et dans l'espace
  • argsort est O(MN log MN) en temps et O(MN) dans l'espace
  • Indexation par sorter est O(MN) dans le temps et dans l'espace
  • Calcul runs est O(MN) dans le temps et O(MN + M) =O(MN) dans l'espace
  • reduceat est une seule passe :O(MN) dans le temps, O(M) dans l'espace
  • Réaffectation A est O(M) dans le temps et dans l'espace

Total :O(MN log MN) temps, O(MN) espace

TL;DR

def make_A(D, I, M):
    index = np.ravel(I)
    data = np.ravel(D)
    sorter = np.argsort(index)
    index = index[sorter]

    if index[0] < 0 or index[-1] >= M:
        raise ValueError('Bad indices')

    data = data[sorter]
    runs = np.r_[0, np.flatnonzero(np.diff(index)) + 1]
    a = np.add.reduceat(data, runs)
    if a.size == M:
        return a
    A = np.zeros(M)
    A[index[runs]] = a
    return A