Find array corresponding to minimal values along an axis in another array

Question

I have two three dimensional arrays, a and b, and want to find the 2D subarray of b with the elements where a had a minimum along the third axis, i.e.

a=n.random.rand(20).reshape((5,2,2))
b=n.arange(20).reshape((5,2,2))
c=n.argmin(a,2) #indices with minimal value of a
d=n.zeros_like(c) #the array I want
for i in range(5):
  for j in range(2):
     d[i,j] = b[i,j,c[i,j]] 

Is there a way I can get these values without the double loop?

I am aware of this answer: replace min value to another in numpy array but if I want this to work for my 3D arrays I'd have to do a lot of reshaping operations - and I'm wondering if there is something simpler.


Show source
| python   | arrays   | numpy   | vectorization   2016-12-26 12:12 3 Answers

Answers ( 3 )

  1. 2016-12-26 12:12

    Here is a Numpythonic way:

    In [83]: x, y, z = a.shape
    In [84]: b[np.repeat(np.arange(x), y), np.tile(np.arange(y), x), c.ravel()].reshape(x, y)
    

    Here np.repeat(np.arange(x), y) will give you the corresponding indices of the first axis.

    In [86]: np.repeat(np.arange(x), y)
    Out[86]: array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4])
    

    np.tile(np.arange(y), x) will give you the corresponding indices of the second axis.

    In [87]: np.tile(np.arange(y), x)
    Out[87]: array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
    

    And for the third one you can just use the flattened shape of c.

    In [88]: c.ravel()
    Out[88]: array([1, 0, 1, 1, 1, 0, 1, 1, 0, 0])
    
  2. 2016-12-26 13:12

    You can use np.ogrid to create a grid for the other dimensions:

    x, y, z = arr.shape  # assuming your array is named "arr"
    xyz = np.ogrid[0:x, 0:y] + [c]  # c is your second axis index (the argmin)
    arr[xyz]
    

    If it's not the last axis then you can simply use insert because ogrid returns a normal python list containing the indices.

  3. 2016-12-26 13:12

    Here's an approach using fancy-indexing -

    m,n,r = b.shape
    d_out = b[np.arange(m)[:,None],np.arange(n),c]
    
◀ Go back