Numpy如何在最大值索引后的切片中查找符合阈值要求的元素索引

实现方案 #

核心思路是生成掩码屏蔽掉最大值索引及之前的位置，将这些位置的差值设为无穷大，再取最小值索引即可，全程为numpy向量化操作，无显式循环：

import numpy as np

# 已有变量data
max_value = np.max(data, axis=0)
threshold_decay = 0.5
decay_value = max_value * threshold_decay
max_indices = np.argmax(data, axis=0)

# 计算所有位置和decay_value的差值
abs_diff = np.abs(data - decay_value)
# 生成axis=0维度的索引数组，形状和data一致
axis0_idx = np.indices(data.shape)[0]
# 将最大值索引及之前的位置差值设为无穷大，避免被选中
abs_diff[axis0_idx <= max_indices] = np.inf
# 取剩余位置差值最小的索引
index_decay = abs_diff.argmin(axis=0)

print(index_decay)

运行后输出为：

array([[4, 4, 4],
       [4, 3, 4],
       [4, 4, 4],
       [4, 4, 4]])

完全符合预期结果，该方案适用于任意规模的三维数组，计算效率远高于循环实现。

内容的提问来源于stack exchange，提问作者Krasto