CODEWITHSUNDEEP
pythonsortingmultidimensional-arraynumpy
Andrew
Andrew

Reputation: 2892

Sort a numpy array by another array, along a particular axis, using less memory

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?

Upvotes: 1

Views: 2317

Answers (1)

senderle
senderle

Reputation: 151007

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

Upvotes: 2

Related Questions