如何向量化实现Pandas DataFrame中当前值前连续较小值行数的高效计算

问题说明 #

需要为Pandas DataFrame的每个值，计算位于它之前、且所有值均小于当前值的最大连续行数。目前用嵌套循环实现了正确结果，但大数据集下速度极慢，需要向量化/高效改写方案。曾参考相关解决方案但结果不符合预期。

原实现代码（嵌套循环） #

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pandas as pd

df_old = pd.DataFrame({'Spam': [10, 1, 2, 3, 4, 5, 6, 8, 3, 4, 7, 8, 9, 1, 2, 7, 11]})
print(df_old, "\n")

for i in range(0, len(df_old)):    
    df = df_old.iloc[0:i+1]
    number_of_rows_with_lower_value = 0

    for j in range(len(df)):
        if (df['Spam'].iloc[-1] > df['Spam'].iloc[-j-1:-1].max()):
            number_of_rows_with_lower_value = j

    print("At row:", i, "value of Spam=", df['Spam'].iloc[-1], "is higher than in last:", number_of_rows_with_lower_value, "rows!")


Spam = df_old['Spam']
Higher=[(Spam[x]>Spam[:x]).sum() for x in range(len(Spam))]
print(Higher)

原代码输出 #

Spam
0     10
1      1
2      2
3      3
4      4
5      5
6      6
7      8
8      3
9      4
10     7
11     8
12     9
13     1
14     2
15     7
16    11 

At row: 0 value of Spam= 10 is higher than in last: 0 rows!
At row: 1 value of Spam= 1 is higher than in last: 0 rows!
At row: 2 value of Spam= 2 is higher than in last: 1 rows!
At row: 3 value of Spam= 3 is higher than in last: 2 rows!
At row: 4 value of Spam= 4 is higher than in last: 3 rows!
At row: 5 value of Spam= 5 is higher than in last: 4 rows!
At row: 6 value of Spam= 6 is higher than in last: 5 rows!
At row: 7 value of Spam= 8 is higher than in last: 6 rows!
At row: 8 value of Spam= 3 is higher than in last: 0 rows!
At row: 9 value of Spam= 4 is higher than in last: 1 rows!
At row: 10 value of Spam= 7 is higher than in last: 2 rows!
At row: 11 value of Spam= 8 is higher than in last: 3 rows!
At row: 12 value of Spam= 9 is higher than in last: 11 rows!
At row: 13 value of Spam= 1 is higher than in last: 0 rows!
At row: 14 value of Spam= 2 is higher than in last: 1 rows!
At row: 15 value of Spam= 7 is higher than in last: 2 rows!
At row: 16 value of Spam= 11 is higher than in last: 16 rows!
[0, 0, 1, 2, 3, 4, 5, 6, 2, 4, 8, 9, 11, 0, 2, 10, 16]

高效解决方案（单调栈法） #

这个问题本质是找每个元素前一个大于等于它的元素的位置，然后用当前索引减去该位置再减1，就是连续小于的最大行数；如果没有这样的元素，结果就是当前索引（因为前面所有元素都小于它）。

用单调栈可以实现O(n)时间复杂度，远快于嵌套循环：

import pandas as pd
import numpy as np

df = pd.DataFrame({'Spam': [10, 1, 2, 3, 4, 5, 6, 8, 3, 4, 7, 8, 9, 1, 2, 7, 11]})
spam = df['Spam'].values

stack = []
result = np.zeros(len(spam), dtype=int)

for i in range(len(spam)):
    # 弹出栈中所有小于当前值的元素
    while stack and spam[stack[-1]] < spam[i]:
        stack.pop()
    # 栈为空说明前面所有元素都小于当前值
    if not stack:
        result[i] = i
    else:
        # 栈顶是前一个大于等于当前值的位置
        result[i] = i - stack[-1] - 1
    stack.append(i)

df['max_consec_lower'] = result

# 输出验证
for idx, row in df.iterrows():
    print(f"At row: {idx} value of Spam= {row['Spam']} is higher than in last: {row['max_consec_lower']} rows!")

print("\nResult array:", result.tolist())

输出结果（与原代码一致） #

At row: 0 value of Spam= 10 is higher than in last: 0 rows!
At row: 1 value of Spam= 1 is higher than in last: 0 rows!
At row: 2 value of Spam= 2 is higher than in last: 1 rows!
At row: 3 value of Spam= 3 is higher than in last: 2 rows!
At row: 4 value of Spam= 4 is higher than in last: 3 rows!
At row: 5 value of Spam= 5 is higher than in last: 4 rows!
At row: 6 value of Spam= 6 is higher than in last: 5 rows!
At row: 7 value of Spam= 8 is higher than in last: 6 rows!
At row: 8 value of Spam= 3 is higher than in last: 0 rows!
At row: 9 value of Spam= 4 is higher than in last: 1 rows!
At row: 10 value of Spam= 7 is higher than in last: 2 rows!
At row: 11 value of Spam= 8 is higher than in last: 3 rows!
At row: 12 value of Spam= 9 is higher than in last: 11 rows!
At row: 13 value of Spam= 1 is higher than in last: 0 rows!
At row: 14 value of Spam= 2 is higher than in last: 1 rows!
At row: 15 value of Spam= 7 is higher than in last: 2 rows!
At row: 16 value of Spam= 11 is higher than in last: 16 rows!

Result array: [0, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 11, 0, 1, 2, 16]

说明 #

• 单调栈维护了一个递减的索引序列，每个元素入栈前会弹出所有比当前值小的元素，确保栈顶是最近的大于等于当前值的元素位置。
• 这种方法避免了嵌套循环的O(n²)时间复杂度，处理大数据集时效率提升显著。

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