Quickly convert numpy arrays with index to dict of numpy arrays keyed on that index

Question

I have a set of numpy arrays. One of these is a list of "keys", and I'd like to rearrange the arrays into a dict of arrays keyed on that key. My current code is:

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

This is pretty slow, since the arrays involved are millions of entries long, and this happens many times. Is it possible to rewrite this in vectorized form? The set of possible keys is known ahead of time, and there are ~10 distinct keys.

Edit: If there are k distinct keys and the list is n long, the current answers are O(nk) (iterate once for each key) and O(n log n) (sort first). I'm still looking for an O(n) vectorized solution, though. This is hopefully possible; after all, the easiest possible nonvectorized thing (i.e. what I already have) is O(n).

Answer 1

Let's import numpy and create some sample data:

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

Now, let's create the dictionary:

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

The idea behind numpy is that C code is much faster than python and numpy provides many common operations coded at the C level. As you mentioned that there were only "~10 distinct keys," that means that the python loop is done only 10 or so times. The rest is C.

Answer 2

The vectorized way to do this is probably going to require you to sort your keys. The basic idea is that you sort the keys and vals to match. Then you can split the val array every time there is a new key in the sorted keys array. The vectorized code looks something like this:

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

This method has complexity O(n * log(n)) because of the argsort step. In practice, argsort is very fast unless n is very large. You're likely going to run into memory issues with this method before argsort gets noticeably slow.

Answer 3

Some 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 used 10 keys, 20 million entries:

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, we see than the sorting technique is a bit faster than letting Numpy find the indices corresponding to each key, but of course both are much much faster than looping in Python. Vectorization is great!

This is on Python 2.7.12, Numpy 1.9.2

Answer 4

defaultdict is intended for building dictionaries like this. In particular is streamlines the step of creating a new dictionary entry for a new key.

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

But with a small known set of keys you don't need this specialized dictionary, since you can easily create the dictionary key entries ahead of time

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

But since you are starting with arrays, array operations to sort and collect like values might well be worth it.