python numpy argmax to multi in multidimensional array

Question

我有以下代码:

import numpy as np
sample = np.random.random((10,10,3))
argmax_indices = np.argmax(sample, axis=2)

即我沿轴= 2取argmax,它给我一个(10,10)矩阵.现在,我想分配这些索引值0.为此,我想索引样本数组.我试过了:

max_values = sample[argmax_indices]

但它不起作用.我想要类似的东西

max_values = sample[argmax_indices]
sample[argmax_indices] = 0

我只是通过检查确定max_values - np.max(sample, axis=2)应该给出零形状矩阵(10,10)来验证.任何帮助将不胜感激.

Answer 1

这是一种方法 -

m,n = sample.shape[:2]
I,J = np.ogrid[:m,:n]
max_values = sample[I,J, argmax_indices]
sample[I,J, argmax_indices] = 0

逐步运行示例

1)样本输入数组:

In [261]: a = np.random.randint(0,9,(2,2,3))

In [262]: a
Out[262]: 
array([[[8, 4, 6],
        [7, 6, 2]],

       [[1, 8, 1],
        [4, 6, 4]]])

2)获取argmax指数axis=2:

In [263]: idx = a.argmax(axis=2)

3)获取用于索引前两个dims的形状和数组:

In [264]: m,n = a.shape[:2]

In [265]: I,J = np.ogrid[:m,:n]

4)使用I,J idx进行索引并使用以下方法存储最大值advanced-indexing:

In [267]: max_values = a[I,J,idx]

In [268]: max_values
Out[268]: 
array([[8, 7],
       [8, 6]])

5)验证我们从zeros减去后得到一个全数组:np.max(a,axis=2)max_values

In [306]: max_values - np.max(a, axis=2)
Out[306]: 
array([[0, 0],
       [0, 0]])

6)再次使用advanced-indexing将这些地点指定为zeros并进行更多级别的可视化验证:

In [269]: a[I,J,idx] = 0

In [270]: a
Out[270]: 
array([[[0, 4, 6], # <=== Compare this against the original version
        [0, 6, 2]],

       [[1, 0, 1],
        [4, 0, 4]]])