Python/Pandas实现：判断数字是否在列表/Series中，不存在则返回下一个最大值

普通列表实现 #

针对普通列表，先确保列表有序，再通过遍历找到符合要求的结果：

nums = [1,2,3,5,7,8,20]

def coerce_num(x, nums_list):
    # 先将列表升序排序，避免无序导致逻辑错误
    sorted_nums = sorted(nums_list)
    if x in sorted_nums:
        return "yes"  # 若需要返回数字本身，可替换为 return x
    # 找到第一个大于x的元素
    for num in sorted_nums:
        if num > x:
            return num
    # 当x大于列表所有元素时，返回最大值
    return sorted_nums[-1]

# 测试用例
print(coerce_num(4, nums))   # 输出: 5
print(coerce_num(9, nums))   # 输出: 20
print(coerce_num(25, nums))  # 输出: 20
print(coerce_num(7, nums))   # 输出: yes

Pandas Series/DataFrame 实现 #

利用Pandas的排序和筛选功能，可高效处理Series或DataFrame列：

Series 处理 #

import pandas as pd

nums_series = pd.Series([1,2,3,5,7,8,20])

def coerce_num_pandas(x, series):
    # 对Series升序排序并重置索引
    sorted_series = series.sort_values().reset_index(drop=True)
    if x in sorted_series.values:
        return "yes"
    # 筛选所有大于x的元素，取第一个（即最小的大于x的数）
    greater_values = sorted_series[sorted_series > x]
    if not greater_values.empty:
        return greater_values.iloc[0]
    # 若x大于所有元素，返回最大值
    return sorted_series.iloc[-1]

# 测试用例
print(coerce_num_pandas(4, nums_series))   # 输出: 5
print(coerce_num_pandas(15, nums_series))  # 输出: 20
print(coerce_num_pandas(30, nums_series))  # 输出: 20

DataFrame 列处理 #

直接将DataFrame的列传入上述函数即可：

df = pd.DataFrame({'values': [1,2,3,5,7,8,20]})
print(coerce_num_pandas(9, df['values']))  # 输出: 20

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