Python如何实现条件不匹配时选取掩码数组的下一个有效值

你当前else分支的逻辑存在缺陷：index_pos_overload[np.nonzero(index_pos_overload)[0][0]]的作用永远是取掩码数组中第一个非0的元素，完全没有和当前传入的possible_overload_cycle做数值比较，因此无论传入的变量值是7还是其他大于5、小于10的数，都会固定返回第一个非零值5，无法匹配到大于当前变量的下一个有效值。

首先可以打印你生成的index_pos_overload确认数组内容：

print(index_pos_overload)
# 输出: [ 0  5 10 15]

该数组默认按升序排列，只需要在变量不存在于数组时，筛选出所有大于possible_overload_cycle的元素，取第一个元素即可得到预期结果。

基础实现（逻辑直观易读） #

用numpy布尔索引做大小筛选，替换原有else分支的逻辑即可：

import numpy as np
import pandas as pd

hyst = pd.DataFrame({"test":[12, 4, 5, 4, 1, 3, 2, 5, 10, 9, 7, 5, 3, 6, 3, 2 ,1, 5, 2]})
possible_overload_cycle = 7
index_pos_overload = np.array([hyst.index[i] for i in range(0, len(hyst)-1, 5)])

if possible_overload_cycle in index_pos_overload:
    hyst_overload_cycle = possible_overload_cycle
else:
    # 筛选所有大于当前变量的元素，取第一个
    larger_values = index_pos_overload[index_pos_overload > possible_overload_cycle]
    # 增加兜底逻辑：如果变量比数组所有值都大，返回数组最后一个有效值，可根据业务调整
    hyst_overload_cycle = larger_values[0] if len(larger_values) > 0 else index_pos_overload[-1]

print(hyst_overload_cycle)

结果验证 #

• 当possible_overload_cycle = 1时，筛选出的大于1的元素为[5,10,15]，取第一个返回5，符合原有正确逻辑
• 当possible_overload_cycle = 7时，筛选出的大于7的元素为[10,15]，取第一个返回10，符合预期
• 当possible_overload_cycle = 12时，筛选出的大于12的元素为[15]，返回15

优化实现（高性能，代码更简洁） #

如果掩码数组长度较大，可以用np.searchsorted做二分查找，时间复杂度更低，且不需要单独判断值是否存在于数组中：

import numpy as np
import pandas as pd

hyst = pd.DataFrame({"test":[12, 4, 5, 4, 1, 3, 2, 5, 10, 9, 7, 5, 3, 6, 3, 2 ,1, 5, 2]})
possible_overload_cycle = 7
index_pos_overload = np.array([hyst.index[i] for i in range(0, len(hyst)-1, 5)])

# 二分查找当前值的插入位置，直接取对应位置的元素即可
insert_pos = np.searchsorted(index_pos_overload, possible_overload_cycle)
hyst_overload_cycle = index_pos_overload[insert_pos] if insert_pos < len(index_pos_overload) else index_pos_overload[-1]

print(hyst_overload_cycle)

该写法在所有测试场景下返回结果和需求完全一致。

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