user3080953 - 3 years ago 102
Python Question

# Problem

I have two numpy arrays,
`A`
and
`indices`
.

`A`
has dimensions m x n x 10000.
`indices`
has dimensions m x n x 5 (output from
`argpartition(A, 5)[:,:,:5]`
).
I would like to get a m x n x 5 array containing the elements of
`A`
corresponding to
`indices`
.

# Attempts

``````indices = np.array([[[5,4,3,2,1],[1,1,1,1,1],[1,1,1,1,1]],
[500,400,300,200,100],[100,100,100,100,100],[100,100,100,100,100]])
A = np.reshape(range(2 * 3 * 10000), (2,3,10000))

A[...,indices] # gives an array of size (2,3,2,3,5). I want a subset of these values
np.take(A, indices) # shape is right, but it flattens the array first
np.choose(indices, A) # fails because of shape mismatch.
``````

# Motivation

I'm trying to get the 5 largest values of
`A[i,j]`
for each
`i<m`
,
`j<n`
in sorted order using
`np.argpartition`
because the arrays can get fairly large.

You can use `advanced-indexing` -

``````m,n = A.shape[:2]
out = A[np.arange(m)[:,None,None],np.arange(n)[:,None],indices]
``````

Sample run -

``````In [330]: A
Out[330]:
array([[[38, 21, 61, 74, 35, 29, 44, 46, 43, 38],
[22, 44, 89, 48, 97, 75, 50, 16, 28, 78],
[72, 90, 48, 88, 64, 30, 62, 89, 46, 20]],

[[81, 57, 18, 71, 43, 40, 57, 14, 89, 15],
[93, 47, 17, 24, 22, 87, 34, 29, 66, 20],
[95, 27, 76, 85, 52, 89, 69, 92, 14, 13]]])

In [331]: indices
Out[331]:
array([[[7, 8, 1],
[7, 4, 7],
[4, 8, 4]],

[[0, 7, 4],
[5, 3, 1],
[1, 4, 0]]])

In [332]: m,n = A.shape[:2]

In [333]: A[np.arange(m)[:,None,None],np.arange(n)[:,None],indices]
Out[333]:
array([[[46, 43, 21],
[16, 97, 16],
[64, 46, 64]],

[[81, 14, 43],
[87, 24, 47],
[27, 52, 95]]])
``````

For getting those indices corresponding to the max 5 elements along the last axis, we would use `argpartition`, like so -

``````indices = np.argpartition(-A,5,axis=-1)[...,:5]
``````

To keep the order from highest to lowest, use `range(5)` instead of `5`.

Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download