Konvertieren Sie schnell numpy Arrays mit Index zu dict von numpy Arrays auf diesen Index

Question

Ich habe eine Reihe von numpy Arrays. Eine davon ist eine Liste von "Schlüsseln", und ich möchte die Arrays in ein Diktat von Arrays umordnen, die auf diesem Schlüssel gekeyed sind. Mein aktueller Code ist:

for key, val1, val2 in itertools.izip(keys, vals1, vals2): 
    dict1[key].append(val1) 
    dict2[key].append(val2) 

Dieses ziemlich langsam ist, da die beteiligten Arrays lang Millionen von Einträgen sind, und dies geschieht oft. Ist es möglich, dies in vektorisierter Form umzuschreiben? Die Menge der möglichen Schlüssel ist vor der Zeit bekannt, und es gibt ~ 10 unterschiedliche Schlüssel.

Edit: Wenn es k verschiedene Schlüssel gibt und die Liste n lang ist, sind die aktuellen Antworten O (nk) (einmal für jeden Schlüssel iterieren) und O (n log n) (zuerst sortieren). Ich suche aber immer noch nach einer O (n) vektorisierten Lösung. Dies ist hoffentlich möglich; Schließlich ist das leicht mögliche nicht-vectorisierte Ding (d. h. was ich bereits habe) O (n).

Answer 1

Die vektorisierte Art, dies zu tun, wird wahrscheinlich erfordern, dass Sie Ihre Schlüssel sortieren. Die Grundidee ist, dass Sie die Schlüssel und Vals entsprechend sortieren. Dann können Sie das val-Array jedes Mal aufteilen, wenn ein neuer Schlüssel im sortierten Schlüssel-Array vorhanden ist. Der vektorisiert Code sieht etwa so aus:

import numpy as np 

keys = np.random.randint(0, 10, size=20) 
vals1 = np.random.random(keys.shape) 
vals2 = np.random.random(keys.shape) 

order = keys.argsort() 
keys_sorted = keys[order] 

# Find uniq keys and key changes 
diff = np.ones(keys_sorted.shape, dtype=bool) 
diff[1:] = keys_sorted[1:] != keys_sorted[:-1] 
key_change = diff.nonzero()[0] 
uniq_keys = keys_sorted[key_change] 

vals1_split = np.split(vals1[order], key_change[1:]) 
dict1 = dict(zip(uniq_keys, vals1_split)) 

vals2_split = np.split(vals2[order], key_change[1:]) 
dict2 = dict(zip(uniq_keys, vals2_split)) 

Diese Methode hat Komplexität O (n * log (n)) wegen des argsort Schritt. In der Praxis ist Argsort sehr schnell, es sei denn, n ist sehr groß. Wahrscheinlich werden Sie mit dieser Methode auf Speicherprobleme stoßen, bevor argsort merklich langsam wird.

Answer 2

Lassen Sie uns Import numpy und erstellen Sie einige Beispieldaten:

>>> import numpy as np 
>>> keys = np.array(('key1', 'key2', 'key3', 'key1', 'key2', 'key1')) 
>>> vals1 = np.arange(6) 
>>> vals2 = np.arange(10, 16) 

Nun lassen Sie uns das Wörterbuch erstellen:

>>> d1 = {}; d2 = {} 
>>> for k in set(keys): 
... d1[k] = vals1[k==keys] 
... d2[k] = vals2[k==keys] 
... 
>>> d1 
{'key3': array([2]), 'key2': array([1, 4]), 'key1': array([0, 3, 5])} 
>>> d2 
{'key3': array([12]), 'key2': array([11, 14]), 'key1': array([10, 13, 15])} 

Die Idee hinter numpy ist, dass C-Code ist viel schneller als Python und numpy bietet viele gemeinsame Operationen, die auf der C-Ebene codiert sind. Da Sie erwähnt haben, dass es nur "~ 10 unterschiedliche Schlüssel" gibt, bedeutet das, dass die Python-Schleife nur etwa 10 Mal ausgeführt wird. Der Rest ist C.

Answer 3

defaultdict ist für den Aufbau solcher Wörterbücher gedacht. Insbesondere wird der Schritt zum Erstellen eines neuen Wörterbucheintrags für einen neuen Schlüssel optimiert.

In [19]: keys = np.random.choice(np.arange(10),100) 
In [20]: vals=np.arange(100) 
In [21]: from collections import defaultdict 
In [22]: dd = defaultdict(list) 
In [23]: for k,v in zip(keys, vals): 
    ...:  dd[k].append(v) 
    ...:  
In [24]: dd 
Out[24]: 
defaultdict(list, 
      {0: [4, 39, 47, 84, 87], 
      1: [0, 25, 41, 46, 55, 58, 74, 77, 89, 92, 95], 
      2: [3, 9, 15, 24, 44, 54, 63, 66, 71, 80, 81], 
      3: [1, 13, 16, 37, 57, 76, 91, 93], 
      ... 
      8: [51, 52, 56, 60, 68, 82, 88, 97, 99], 
      9: [21, 29, 30, 34, 35, 59, 73, 86]}) 

Aber mit einem kleinen bekannten Satz von Schlüsseln Sie dieses Fachwörterbuch nicht benötigen, da Sie leicht die Wörterbuchschlüsseleinträge vor der Zeit

dd = {k:[] for k in np.unique(keys)} 

erstellen Aber da Sie mit Arrays beginnen, werden , Array-Operationen zum Sortieren und Sammeln ähnlicher Werte können sich durchaus lohnen.

Answer 4

Einige Timings:

import numpy as np 
import itertools 

def john1024(keys, v1, v2): 
    d1 = {}; d2 = {}; 
    for k in set(keys): 
    d1[k] = v1[k==keys] 
    d2[k] = v2[k==keys] 
    return d1,d2 

def birico(keys, v1, v2): 
    order = keys.argsort() 
    keys_sorted = keys[order] 
    diff = np.ones(keys_sorted.shape, dtype=bool) 
    diff[1:] = keys_sorted[1:] != keys_sorted[:-1] 
    key_change = diff.nonzero()[0] 
    uniq_keys = keys_sorted[key_change] 
    v1_split = np.split(v1[order], key_change[1:]) 
    d1 = dict(zip(uniq_keys, v1_split)) 
    v2_split = np.split(v2[order], key_change[1:]) 
    d2 = dict(zip(uniq_keys, v2_split)) 
    return d1,d2 

def knzhou(keys, v1, v2): 
    d1 = {k:[] for k in np.unique(keys)} 
    d2 = {k:[] for k in np.unique(keys)} 
    for key, val1, val2 in itertools.izip(keys, v1, v2): 
    d1[key].append(val1) 
    d2[key].append(val2) 
    return d1,d2 

I 10 Tasten, 20 Millionen Einträge:

import timeit 

keys = np.random.randint(0, 10, size=20000000) #10 keys, 20M entries 
vals1 = np.random.random(keys.shape) 
vals2 = np.random.random(keys.shape) 

timeit.timeit("john1024(keys, vals1, vals2)", "from __main__ import john1024, keys, vals1, vals2", number=3) 
11.121668815612793 
timeit.timeit("birico(keys, vals1, vals2)", "from __main__ import birico, keys, vals1, vals2", number=3) 
8.107877969741821 
timeit.timeit("knzhou(keys, vals1, vals2)", "from __main__ import knzhou, keys, vals1, vals2", number=3) 
51.76217794418335 

So sehen wir als die Sortiertechnik ein bisschen schneller als im Stich gelassen Numpy die Indizes zu entsprechenden finden ist jeder Schlüssel, aber natürlich sind beide viel viel schneller als das Schleifen in Python. Vektorisierung ist großartig!

Dies ist auf Python 2.7.12, Numpy 1.9.2