Numpy中按索引数组求对应最大值并替换原数组的向量化实现

给定数组 a = [0.51, 0.6, 0.8, 0.65, 0.7, 0.75, 0.9] 和索引数组 ind = [0, 1, 3, 1, 2, 2, 3]，需要为ind中的每个索引值，找到原数组a中对应位置的最大值，并用该最大值替换原数组中对应索引位置的元素，预期输出为 out = [0.51, 0.65, 0.9, 0.65, 0.75, 0.75, 0.9]。

举个例子：ind中的值1对应a里的元素[0.6, 0.65]，最大值是0.65，需要替换a中位置1和3的元素。

以下是几种无显式循环的向量化实现方案：

方案一：利用np.unique的逆映射快速分组 #

这种写法简洁直观，适合大多数场景：

import numpy as np

a = np.array([0.51, 0.6, 0.8, 0.65, 0.7, 0.75, 0.9])
ind = np.array([0, 1, 3, 1, 2, 2, 3])

# 获取唯一索引，以及原ind中每个元素对应唯一索引的位置
unique_ind, inv = np.unique(ind, return_inverse=True)
# 计算每个唯一索引对应的a中元素的最大值
group_max = np.array([a[ind == idx].max() for idx in unique_ind])
# 通过逆映射将最大值映射回原数组长度
out = group_max[inv]
print(out)
# 输出：[0.51 0.65 0.9  0.65 0.75 0.75 0.9 ]

方案二：用np.argsort+np.maximum.reduceat实现纯向量化 #

如果想要完全避免列表推导，可通过排序+分段取最大值的方式实现：

import numpy as np

a = np.array([0.51, 0.6, 0.8, 0.65, 0.7, 0.75, 0.9])
ind = np.array([0, 1, 3, 1, 2, 2, 3])

# 按ind对原数组索引排序
sort_idx = np.argsort(ind)
sorted_ind = ind[sort_idx]
sorted_a = a[sort_idx]

# 找到分组的边界（每个新组的起始位置）
group_boundaries = np.r_[True, sorted_ind[1:] != sorted_ind[:-1]]
split_points = np.where(group_boundaries)[0]

# 用reduceat分段计算每组的最大值
group_max = np.maximum.reduceat(sorted_a, split_points)

# 生成与原数组长度匹配的结果，再恢复原顺序
out = np.empty_like(a)
out[sort_idx] = np.repeat(group_max, np.diff(np.r_[split_points, len(sorted_ind)]))
print(out)
# 输出：[0.51 0.65 0.9  0.65 0.75 0.75 0.9 ]

原循环实现代码 #

你提供的循环版本代码如下：

import numpy as np

a = np.array([0.51, 0.6, 0.8, 0.65, 0.7, 0.75, 0.9])
ind = np.array([0, 1, 3, 1, 2, 2, 3])

# 获取唯一索引
unique_indices = np.unique(ind)

# 遍历唯一索引，找到每个索引对应的最大值并替换
for index in unique_indices:
    max_value = np.max(a[ind == index])
    a[ind == index] = max_value

out = a
print(out)

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