1

From the answer to this question, I learned how to sort the entries of one numpy array a by the values of another numpy array b, along a particular axis.

However, this method requires the creation of several intermediate arrays that are the same size as a, one for each dimension of a. Some of my arrays are quite large, and this becomes inconvenient. Is there a way to accomplish the same goal that uses less memory?

Community
  • 1
  • 1
Andrew
  • 2,842
  • 5
  • 31
  • 49
  • Will the array `a` in the linked question always consist of blocks of equal numbers? In this case, it would be rather easy to give a better solution. – Sven Marnach May 28 '11 at 23:03
  • Good question, Sven. In my real code, the array to be sorted is effectively random, I just gave a fairly trivial example for illustration. – Andrew May 29 '11 at 12:38

1 Answers1

2

Would a record array serve your purposes?

>>> a = numpy.zeros((3, 3, 3))
>>> a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
>>> b = numpy.arange(3*3*3).reshape((3, 3, 3))
>>> c = numpy.array(zip(a.flatten(), b.flatten()), dtype=[('f', float), ('i', int)]).reshape(3, 3, 3)
>>> c.sort(axis=0)
>>> c['i']
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]]])

A cleaner way to generate the coupled array:

>>> c = numpy.rec.fromarrays([a, b], dtype=[('f', float), ('i', int)])

or

>>> c = numpy.rec.fromarrays([a, b], names='f, i')
senderle
  • 145,869
  • 36
  • 209
  • 233
  • I should add that this is very similar to [this question](http://stackoverflow.com/questions/1903462/how-can-i-zip-sort-parallel-numpy-arrays) but since that one doesn't address multidimensional arrays explicitly, I think it's worth having a separate question. – senderle May 28 '11 at 14:48
  • Neat, I didn't know about record arrays. I'll try this out. – Andrew May 29 '11 at 12:38
  • The record array approach was exactly what I was looking for. Great answer. – Andrew May 30 '11 at 13:42
  • It's great that you can specify which field of the record array to sort by with the 'order' keyword argument, too. – Andrew May 30 '11 at 13:44