如何用Numpy统计数组中元素后续跌幅超100的符合条件数量

核心思路 #

要高效实现需求，不需要嵌套遍历每个元素，只需要计算每个位置之后所有元素的最小值，再判断当前元素减去后续最小值是否≥100即可：只要后续最小值满足下降幅度要求，就证明存在符合条件的元素，反之则不存在。

完整代码 #

import numpy as np

def count_valid_starts(arr):
    # 计算从右向左的累积最小值，得到每个位置到数组末尾的最小值
    min_suffix = np.minimum.accumulate(arr[::-1])[::-1]
    # 构造每个位置之后的最小值数组，最后一位补无穷大（最后一个元素无后续元素，不可能满足条件）
    min_after = np.append(min_suffix[1:], np.inf)
    # 统计所有满足下降幅度要求的起始点数量
    return np.sum(arr - min_after >= 100)

# 测试用例
val = np.array([7.4, 1180.9, 0.6, 103.0, 58.8, 35, -47, 47.2, 78.1, 37.8])
val2 = np.array([46.5, 55.7, 7.0, 19.6, 7.6, 36.5, 34.7, 101.9, 179.7, 85.5])
val3 = np.array([120, 20, -80, -5.5])

print(count_valid_starts(val))
print(count_valid_starts(val2))
print(count_valid_starts(val3))

输出结果 #

3
0
2

补充说明 #

你之前尝试用np.subtract.accumulate得不到正确结果，是因为该方法实现的是累积减法逻辑，和我们需要求后续最小值的需求不匹配，换成np.minimum.accumulate做后缀最小值计算即可满足要求。

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