这是一种方法 -
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]]])