Sort an ND numpy array by another 1-D array

Question

From the answer to this question ( Sorting a numpy array by another array along a specific axis using less memory ), I learned how to sort a multidimensional numpy array with the avalues ​​of another numpy array bwithout creating too many additional arrays.

However, numpy.rec.fromarrays([a, b])only works if the arrays aand bhave the same shape. My array bis a 1-D array, but the array ais an ND array (N is not specified). Is it good (and efficient) to sort an array aamong a particular axis by the value of a 1-D array b?

Answer 1

Use np.takewith keyword argument axis:

>>> a = np.arange(2*3*4).reshape(2, 3, 4)
>>> a
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])
>>> b = np.arange(3)
>>> np.random.shuffle(b)
>>> b
array([1, 0, 2])
>>> np.take(a, b, axis=1)
array([[[ 4,  5,  6,  7],
        [ 0,  1,  2,  3],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [12, 13, 14, 15],
        [20, 21, 22, 23]]])

If you want to use fancy indexing, you just need to impose an indexing tuple on sufficiently empty fragments:

>>> a[:, b]
array([[[ 4,  5,  6,  7],
        [ 0,  1,  2,  3],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [12, 13, 14, 15],
        [20, 21, 22, 23]]])

Or in a more general setup:

>>> axis = 1
>>> idx = (slice(None),) * axis + (b,)
>>> a[idx]
array([[[ 4,  5,  6,  7],
        [ 0,  1,  2,  3],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [12, 13, 14, 15],
        [20, 21, 22, 23]]])

But np.takereally should be your first option.