如何在Numpy中对np.where返回的连续True索引分组求最大值？

问题描述 #

编写Numpy脚本时，需要针对np.where返回的连续True索引区间，给每个区间内的longband元素赋值为该区间的最大值。现有代码：

RsiMa = ...array calculations...
RsiMa_1 = np.roll(RsiMa, 1)
longband = ...array calculations...
longband_1 = np.roll(longband, 1)

long_max = np.where(((RsiMa_1 > longband_1) & (RsiMa > longband_1)))

相关数组数值：

期望执行赋值后，每个连续True区间的元素被替换为区间内最大值：

但直接执行longband[long_max] = np.max(longband[long_max])会把所有True位置替换为全局最大值52，不符合需求。需要不使用循环实现目标。

实现方案 #

步骤1：转换为布尔掩码 #

先把np.where返回的索引转换成布尔数组，方便后续处理：

# 将索引转为布尔掩码
bool_mask = np.zeros_like(longband, dtype=bool)
bool_mask[long_max] = True

步骤2：识别连续区间的起止位置 #

通过计算掩码的差分，找出每个连续True区间的起始和结束索引：

# 计算掩码的变化点：1表示从False转True（区间开始），-1表示从True转False（区间结束）
diff = np.diff(bool_mask.astype(int), prepend=0, append=0)
starts = np.where(diff == 1)[0]
ends = np.where(diff == -1)[0]

步骤3：向量化赋值（无显式循环） #

计算每个区间的最大值，再通过np.repeat生成对应长度的最大值数组，最后赋值：

# 计算每个连续区间的最大值
seg_maxes = np.array([np.max(longband[s:e]) for s, e in zip(starts, ends)])
# 生成与掩码匹配的最大值数组
max_vals = np.repeat(seg_maxes, ends - starts)
# 给对应位置赋值
longband[bool_mask] = max_vals

如果偏好显式循环（代码更直观，数据量不大时效率无差异），也可以用：

for start, end in zip(starts, ends):
    seg_max = np.max(longband[start:end])
    longband[start:end] = seg_max

执行后即可得到期望的结果，每个连续True区间的元素被替换为该区间的局部最大值。

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