I know there must be an elegant way to do this using advanced indexing, I just can't figure it out.
Suppose I have the (2,3,4) array
x = 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]]])
and the (4,) array y = array([1,0,1,1])
What is the most elegant way to obtain the (3,4) array that
z = np.zeros((3,4))
for ii in xrange(3):
for jj in xrange(4):
z[ii,jj] = x[y[jj],ii,jj]
produces?
In [490]: x[y,:,np.arange(4)]
Out[490]:
array([[12, 16, 20],
[ 1, 5, 9],
[14, 18, 22],
[15, 19, 23]])
We need to transpose this. With a mix of basic and advanced indexing, the slice dimension has been put last:
In [491]: x[y,:,np.arange(4)].T
Out[491]:
array([[12, 1, 14, 15],
[16, 5, 18, 19],
[20, 9, 22, 23]])
(that basic/advanced quirk is documented and discussed in some SO.)
or with advanced indexing all around:
In [492]: x[y,np.arange(3)[:,None],np.arange(4)]
Out[492]:
array([[12, 1, 14, 15],
[16, 5, 18, 19],
[20, 9, 22, 23]])