arraysnumpyadvanced-indexing

Easy way to do nd-array contraction using advanced indexing in Python


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?


Solution

  • 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]])