Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Welcome To Ask or Share your Answers For Others

Categories

0 votes
433 views
in Technique[技术] by (71.8m points)

python - Sort N-D numpy array by another 1-D array

From the answer to this question (Sort a numpy array by another array, along a particular axis, using less memory), I learned how to sort a multidimensional numpy array a by the values of another numpy array b without creating too many additional arrays.

However, numpy.rec.fromarrays([a, b]) only works if the arrays a and b have the same shape. My b array is 1-D array, but a array is a N-D array (N is not specified). Is it a nice way (and efficient) to sort the a array among a particular axis by the value of the 1-D array b?

See Question&Answers more detail:os

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
Welcome To Ask or Share your Answers For Others

1 Reply

0 votes
by (71.8m points)

Use np.take with the axis keyword argument:

>>> 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 pad the indexing tuple with enough empty slices:

>>> 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 setting:

>>> 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.take should really be your first option.


与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
OGeek|极客中国-欢迎来到极客的世界,一个免费开放的程序员编程交流平台!开放,进步,分享!让技术改变生活,让极客改变未来! Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Click Here to Ask a Question

...