Question

I have a question regarding how to extract certain values from a 2D numpy array

Foo = 
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

Bar =
array([[0, 0, 1],
       [1, 2, 3]])

I want to extract elements from Foo using the values of Bar as indices, such that I end up with an 2D matrix/array Baz of the same shape as Bar. The ith column in Baz correspond is Foo[(np.array(each j in Bar[:,i]),np.array(i,i,i,i ...))]

Baz =
array([[ 1,  2,  6],
       [ 4,  8, 12]])

I could do a couple nested for-loops but I was wondering if there is a more elegant, numpy-ish way to do this.

Sorry if this is a bit convoluted. Let me know if I need to explain further.

Thanks!

Answer 1

You can use Bar as the row index and an array [0, 1, 2] as the column index:

# for easy copy-pasting
import numpy as np
Foo = np.array([[ 1,  2,  3], [ 4,  5,  6], [ 7,  8,  9],  [10, 11, 12]])
Bar = np.array([[0, 0, 1], [1, 2, 3]])

# now use Bar as the `i` coordinate and 0, 1, 2 as the `j` coordinate:

Foo[Bar, [0, 1, 2]]
# array([[ 1,  2,  6],
#        [ 4,  8, 12]])

# OR, to automatically generate the [0, 1, 2]

Foo[Bar, xrange(Bar.shape[1])]