使用NumPy和Pandas查找数值序列中的局部最大值（峰值索引）

需求说明 #

从输入的空格分隔整数列表创建Pandas Series，找出所有局部最大值（峰值）的索引。峰值定义为两侧数值均小于自身的元素，输出格式为索引数组。

• 样例输入：12 1 2 1 9 10 2 5 7 8 9 -9 10 5 15
• 样例输出：[2 5 10 12]

实现代码 #

import numpy as np
import pandas as pd

# 读取输入并转换为整数序列
input_nums = list(map(int, input().strip().split()))
series = pd.Series(input_nums)

# 生成峰值判断条件
# 当前元素 > 前一个元素
greater_than_prev = series > series.shift(1)
# 当前元素 > 后一个元素
greater_than_next = series > series.shift(-1)
# 同时满足两个条件的位置即为峰值
peak_mask = greater_than_prev & greater_than_next

# 提取峰值索引并转为NumPy数组输出
peak_indices = series[peak_mask].index.to_numpy()
print(peak_indices)

代码说明 #

1. 输入处理：读取控制台输入，拆分并转换为整数列表，再创建Pandas Series。
2. 峰值判断：
   • series.shift(1)：将序列元素向后偏移一位，用于比较当前元素与前一个元素的大小。
   • series.shift(-1)：将序列元素向前偏移一位，用于比较当前元素与后一个元素的大小。
   • 两个比较结果取逻辑与，得到仅满足“大于两侧元素”的位置掩码（首尾元素因偏移后产生NaN，会自动被排除）。
3. 结果输出：通过掩码筛选出峰值对应的索引，转为NumPy数组后打印。

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