4

I understand that through Sort N-D numpy array by another 1-D array the use of fancy indexing, I can do the following c = a[:, :, b] with b defining the order to which I want to sort by column

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

       [[ 5,  4],
        [ 7,  6]],

       [[ 9,  8],
        [11, 10]]])

Now I increase b with 2 more inputs to b2 corresponding to how I want to sort each set of 2x2 in a 

>>> b2 = np.array([[1, 0], [0, 1], [1, 0]])
>>> c2 = ?
>>> c2
array([[[ 1,  0],
        [ 3,  2]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 9,  8],
        [11, 10]]])

I have a larger set of inputs and I have a function that returns an array similar to 'b2' which provides me the info to which I should obtain. Hence may I know what should I filling into the c2 = ? in order to get the desired result?

Improve this question
2
  • Shouldn't b2 be np.array([[1, 0], [0, 1], [1, 0]]) (last row swapped) for that expected output? Commented Dec 7, 2015 at 12:12
  • @ajcr Thanks for spotting that mistake! :) Commented Dec 7, 2015 at 12:13

3 Answers 3

Reset to default
5

Here's one approach with fancy-indexing -

(a[np.arange(a.shape[0])[:,None],:,b2]).transpose(0,2,1)

Sample run -

In [191]: a
Out[191]: 
array([[[7, 8, 5, 2, 0],
        [6, 7, 0, 7, 1],
        [7, 6, 5, 4, 0]],

       [[8, 0, 5, 5, 7],
        [4, 3, 4, 0, 1],
        [8, 6, 3, 2, 4]],

       [[3, 2, 7, 3, 7],
        [4, 3, 0, 1, 5],
        [4, 3, 7, 8, 7]]])

In [192]: b2
Out[192]: 
array([[1, 2, 4, 3, 0],
       [4, 2, 0, 1, 3],
       [1, 3, 4, 0, 2]])

In [193]: (a[np.arange(a.shape[0])[:,None],:,b2]).transpose(0,2,1)
Out[193]: 
array([[[8, 5, 0, 2, 7],
        [7, 0, 1, 7, 6],
        [6, 5, 0, 4, 7]],

       [[7, 5, 8, 0, 5],
        [1, 4, 4, 3, 0],
        [4, 3, 8, 6, 2]],

       [[2, 3, 7, 3, 7],
        [3, 1, 5, 4, 0],
        [3, 8, 7, 4, 7]]])
Improve this answer
Sign up to request clarification or add additional context in comments.

2 Comments

This is a nice example of mixing advanced and basic indexing.
@unutbu And I didn't even know that I was mixing things up! :)
1

In case nobody finds a pure fancy-indexing solution, this here is one that loops through the first axis:

np.asarray([a[n,:,p] for n,p in enumerate(b2)])

array([[[ 1,  3],
    [ 0,  2]],

   [[ 4,  6],
    [ 5,  7]],

   [[ 9, 11],
    [ 8, 10]]])
Improve this answer

1 Comment

Thanks for your answer! A fancy-indexing solution will be better bcos speed is what I'm looking for. :)
1

Something similar to @Divakar's solution but without the transpose.

In [259]: I,J, K = np.ogrid[:3,:2,:2]

In [260]: a[I, J, b[:,None,:]]
Out[260]: 
array([[[ 1,  0],
        [ 3,  2]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 9,  8],
        [11, 10]]])

I'm using ogrid (or np.ix_) as a compact way of generating 2 3d arrays that broadcast with b[:,None,:] to produce a set of (3,2,2) indices.

The equivalent with full Nones is:

a[np.arange(3)[:,None,None], np.arange(2)[None,:,None], b[:,None,:]]

(this makes it clearer that b is selecting items along the 1st and last axes)

To see the full broadcasted indexing arrays, print:

 np.broadcast_arrays(I,J,b[:,None,:])
Improve this answer

1 Comment

Thanks! Similar to Divakar's answer, I learnt a lot and thanks for also explaining it step by step.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.