How to place elements of a NumPy array in a desired order using 2D array of indices

Question

I have a NumPy array of the form

a = [[[a1, a2, a3],
      [a4, a5, a6],
      [a7, a8, a9]],

     [[b1, b2, b3],
      [b4, b5, b6],
      [b7, b8, b9]]]

and another array of the form

M = [[1, 0, 2],
     [1, 2, 0]]

Is there an easy way to reorder the first array using the indices in the second one? So in the end the first one would become:

a = [[[a4, a5, a6],
      [a1, a2, a3],
      [a7, a8, a9]],

     [[b4, b5, b6],
      [b7, b8, b9],
      [b1, b2, b3]]]

According to the duplicate target your example output can be achieved by `a[np.arange(2)[:,None], M]`, where both `a` and `M` are NumPy arrays. — Georgy, Oct 30 '19 at 13:39

rassar · Answer 1 · 2019-10-30T13:57:10.017

0

You can use a list comprehension -

b = [[a[i][j] for j in m] for i, m in enumerate(M)]

You can also use this:

b = [i[j] for i, j in zip(a, M)]

since your arrays are numpy arrays, but it will be slower. The list comprehension is plenty fast for any purpose.

If you don't care about loading the whole list at a time, you can use a generator:

b = ([a[i][j] for j in m] for i, m in enumerate(M))

but this can only be iterated over once.

You can also use the numpy methods:

b = a[np.arange(len(M))[:,None],M]

which will be faster than the list comprehension for larger arrays.

edited Oct 30 '19 at 13:57

answered Oct 30 '19 at 13:24

rassar

Thank you for your reply! Is this fast? My real a array is 700,000 dimensional, not just 2 dimensional as in this example. Would that for loop be good enough? – JohnDoe122 Oct 30 '19 at 13:27
It's a list comprehension, so it should be extremely fast. – rassar Oct 30 '19 at 13:33
1

"_But this is much slower than the list comprehension_" - for small arrays this might be true, but check the times for arrays, like, 100 times larger. – Georgy Oct 30 '19 at 13:50
@Georgy You're right, it is much faster. – rassar Oct 30 '19 at 13:56

1 Answers1