用Pandas向量化方法查找列中首个后续小于当前值的索引

一、计算F(Val)：每行之后首个小于当前值的索引 #

针对大数据集，逐行循环的O(n²)复杂度会导致严重性能问题，单调栈是线性时间O(n)的最优解法，完全避免冗余计算：

实现代码 #

import numpy as np
import pandas as pd

df = pd.DataFrame({'Val': [1, 2, 3, 4, 0, 1, -1, -2, -3]}, index=np.arange(0,9))

# 初始化单调栈与结果数组
stack = []
f_val = np.full(len(df), np.nan)

# 从后往前遍历，维护单调递增栈
for i in reversed(range(len(df))):
    # 弹出栈中所有不小于当前值的元素，保证栈顶是下一个更小元素
    while stack and df['Val'].iloc[stack[-1]] >= df['Val'].iloc[i]:
        stack.pop()
    # 栈不为空时，栈顶即为目标索引
    if stack:
        f_val[i] = stack[-1]
    stack.append(i)

df['F(Val)'] = f_val

核心逻辑 #

单调栈始终维护一个对应值单调递增的索引序列：

• 从后往前遍历，每次清理栈中不小于当前值的元素，剩余栈顶就是第一个比当前值小的元素索引
• 每个元素仅入栈、出栈各一次，时间复杂度严格O(n)，处理百万级数据也能保持高效

二、计算G(Val)：当前索引到F(Val)区间内的最大值 #

利用F(Val)的单调性质，我们可以通过反向递推实现O(n)时间的高效计算，无需重复遍历区间：

实现代码 #

g_val = np.full(len(df), np.nan)

# 从倒数第二行反向递推
for i in reversed(range(len(df)-1)):
    # 若下一个更小元素就是相邻行，区间最大值就是当前值
    if df['F(Val)'].iloc[i] == i+1:
        g_val[i] = df['Val'].iloc[i]
    else:
        # 否则区间最大值为当前值与下一行G(Val)的较大值
        g_val[i] = max(df['Val'].iloc[i], g_val[i+1])

df['G(Val)'] = g_val

核心逻辑 #

观察区间规律：[i, F(Val)]的最大值，等于当前值Val[i]和[i+1, F(Val)]的最大值（即G(Val)[i+1]）中的较大者，而F(Val)的单调性保证了[i+1, F(Val)]的最大值已被提前计算。

最终结果 #

运行代码后，DataFrame将生成期望的结果列：

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