Найти массив, соответствующий минимальным значениям вдоль оси в другом массиве

Question

У меня есть два трехмерных массива a и b и вы хотите найти 2D-подрамник b с элементами, где a имел минимум вдоль третьей оси, то есть

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

Есть ли способ получить эти значения без двойной петли?

Я знаю этот ответ: replace min value to another in numpy array , но если я хочу, чтобы это работало для моих 3D массивов я должен был бы сделать много операций перестройки - и мне интересно, если есть что-то проще.

Answer 1

Вот подход, использующий fancy-indexing -

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

Answer 2

Вот Numpythonic путь:

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) 

Здесь np.repeat(np.arange(x), y) даст вам соответствующие показатели первой оси.

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) даст вам соответствующие показатели второй оси.

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

А для третьего вы можете просто использовать уплощенную форму c.

In [88]: c.ravel() 
Out[88]: array([1, 0, 1, 1, 1, 0, 1, 1, 0, 0]) 

Answer 3

Вы можете использовать np.ogrid для создания сетки для других размеров:

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] 

Если это не последняя ось, то вы можете просто использовать insert, потому что ogrid возвращает нормальный список питона, содержащий индексы.