CODEWITHSUNDEEP
numpy
Martin Wang
Martin Wang

Reputation: 1007

Reduce a dimension of numpy array by selecting

I have a 3d array

A = np.random.random((4,4,3))

and a index matrix

B = np.int_(np.random.random((4,4))*3)

How do I get a 2D array from A based on index matrix B?

In general, how to get a N-1 dimensional array from a ND array and a N-1 dimensional index array?

Upvotes: 4

Views: 3270

Answers (2)

butterfly
butterfly

Reputation: 47

If you prefer to use mesh as suggested by Daniel, you may also use

A[tuple( np.ogrid[:A.shape[0], :A.shape[1]] + [B] )]

to work with sparse indices. In the general case you could use

A[tuple( np.ogrid[ [slice(0, end) for end in A.shape[:-1]] ] + [B] )]

Note that this may also be used when you'd like to index by B an axis different from the last one (see for example this answer about inserting an element into a list).

Otherwise you can do it using broadcasting:

A[np.arange(A.shape[0])[:, np.newaxis], np.arange(A.shape[1])[np.newaxis, :], B]

This may be generalized too but it's a bit more complicated.

Upvotes: 0

Daniel
Daniel

Reputation: 19547

Lets take an example:

>>> A = np.random.randint(0,10,(3,3,2))
>>> A
array([[[0, 1],
        [8, 2],
        [6, 4]],

       [[1, 0],
        [6, 9],
        [7, 7]],

       [[1, 2],
        [2, 2],
        [9, 7]]])

Use fancy indexing to take simple indices. Note that the all indices must be of the same shape and the shape of each index will be what is returned.

>>> ind = np.arange(2)
>>> A[ind,ind,ind]
array([0, 9]) #Index (0,0,0) and (1,1,1)

>>> ind = np.arange(2).reshape(2,1)
>>> A[ind,ind,ind]
array([[0],
       [9]])

So for your example we need to supply the grid for the first two dimensions:

>>> A = np.random.random((4,4,3))
>>> B = np.int_(np.random.random((4,4))*3)
>>> A
array([[[ 0.95158697,  0.37643036,  0.29175815],
        [ 0.84093397,  0.53453123,  0.64183715],
        [ 0.31189496,  0.06281937,  0.10008886],
        [ 0.79784114,  0.26428462,  0.87899921]],

       [[ 0.04498205,  0.63823379,  0.48130828],
        [ 0.93302194,  0.91964805,  0.05975115],
        [ 0.55686047,  0.02692168,  0.31065731],
        [ 0.92822499,  0.74771321,  0.03055592]],

       [[ 0.24849139,  0.42819062,  0.14640117],
        [ 0.92420031,  0.87483486,  0.51313695],
        [ 0.68414428,  0.86867423,  0.96176415],
        [ 0.98072548,  0.16939697,  0.19117458]],

       [[ 0.71009607,  0.23057644,  0.80725518],
        [ 0.01932983,  0.36680718,  0.46692839],
        [ 0.51729835,  0.16073775,  0.77768313],
        [ 0.8591955 ,  0.81561797,  0.90633695]]])
>>> B
array([[1, 2, 0, 0],
       [1, 2, 0, 1],
       [2, 1, 1, 1],
       [1, 2, 1, 2]])

>>> x,y = np.meshgrid(np.arange(A.shape[0]),np.arange(A.shape[1]))
>>> x
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])
>>> y
array([[0, 0, 0, 0],
       [1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3]])

>>> A[x,y,B]
array([[ 0.37643036,  0.48130828,  0.24849139,  0.71009607],
       [ 0.53453123,  0.05975115,  0.92420031,  0.36680718],
       [ 0.10008886,  0.02692168,  0.86867423,  0.16073775],
       [ 0.26428462,  0.03055592,  0.16939697,  0.90633695]])

Upvotes: 4

Related Questions