在Pandas DataFrame中实现当前值与前后n个值的条件比较

实现思路 #

利用Pandas的shift函数生成当前值前后n个位置的偏移列，逐行统计满足指定条件的有效次数，最后根据次数是否≥q来标记1或0。

示例代码 #

import pandas as pd

# 初始化示例数据
df = pd.DataFrame({'x': [2, 2, 2, 1, 1, 2, 2]})

# 参数配置
n = 2  # 前后各n个值
q = 3  # 满足条件的最小次数
# 定义比较条件：当前值 <= 对比值/2
def compare(current, val):
    return current <= val / 2

# 生成所有需要对比的偏移列
shifted_columns = []
# 添加向前偏移1~n位的列（当前值的前n个值）
for i in range(1, n+1):
    shifted_columns.append(df['x'].shift(i))
# 添加向后偏移1~n位的列（当前值的后n个值）
for i in range(1, n+1):
    shifted_columns.append(df['x'].shift(-i))

# 合并偏移列为临时DataFrame
shifted_df = pd.concat(shifted_columns, axis=1)

# 统计每行满足条件的有效次数（排除NaN，即边界无值的情况）
match_counts = shifted_df.apply(
    lambda row: sum(compare(df['x'].iloc[row.name], val) for val in row if pd.notna(val)),
    axis=1
)

# 生成最终标记列
df['comparison'] = (match_counts >= q).astype(int)

print(df)

代码说明 #

1. 偏移列生成：通过循环生成向前、向后各n个偏移列，对应当前值前后n个位置的数据。
2. 条件统计：逐行遍历偏移列的有效值（跳过边界的NaN），用自定义条件判断并统计满足次数。
3. 标记生成：将满足次数≥q的行标记为1，否则为0，转为整数类型。

输出结果 #

x  comparison
0  2           0
1  2           0
2  2           0
3  1           1
4  1           1
5  2           0
6  2           0

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