Pandas高效统计当前行前后连续满足特定条件的行数的方案需求

嘿，针对你25k行DataFrame的这个效率问题，我来给你几个更高效的解决方案！原apply方法因为每次都要切片、反转DataFrame，产生大量临时对象，确实在数据量大的时候会很慢。我们可以通过预计算+高效循环，甚至用numba加速来大幅提升速度。

先明确问题核心 #

我们需要给每一行统计：

• next count：从当前行的下一行开始，连续满足「该行A > 当前行A 且 该行A < 该行B」的行数，直到第一个不满足的行就停止计数
• previous count：从当前行的上一行往前倒序，连续满足相同条件的行数，同样到第一个不满足的行停止

原方法的瓶颈在于每次行级apply都要做DataFrame切片和反转，这在25k行时会产生巨大的性能开销。

优化方案1：用Numpy数组替代DataFrame操作 #

先预计算好基础条件，再用纯数组循环，避免DataFrame的额外开销：

import pandas as pd
import numpy as np

# 你的示例数据
df = pd.DataFrame({'A' : [0,3,2,9,1,0,4,7,3,2], 'B': [9,8,3,5,5,5,5,8,0,4]})

# 预计算每个行是否满足 A < B（这是后续/前序行需要满足的基础条件之一）
mask = (df['A'] < df['B']).values
A_vals = df['A'].values
n = len(df)

# 计算next count
next_count = np.zeros(n, dtype=int)
for i in range(n):
    cnt = 0
    # 从下一行开始遍历，直到不满足条件
    for j in range(i+1, n):
        if mask[j] and A_vals[j] > A_vals[i]:
            cnt += 1
        else:
            break
    next_count[i] = cnt

# 计算previous count
prev_count = np.zeros(n, dtype=int)
for i in range(n):
    cnt = 0
    # 从上一行往前倒序遍历
    for j in range(i-1, -1, -1):
        if mask[j] and A_vals[j] > A_vals[i]:
            cnt += 1
        else:
            break
    prev_count[i] = cnt

# 合并结果到原DataFrame
df['next count'] = next_count
df['previous count'] = prev_count

这个方法比原apply快很多，因为用Numpy数组直接遍历，避免了DataFrame切片带来的内存和时间开销。

优化方案2：用Numba加速循环（性能提升几十倍） #

如果25k行还是觉得不够快，用numba对循环进行JIT编译，能把纯Python循环的速度提升到接近C的水平：

from numba import jit

# 用numba编译next count的计算函数
@jit(nopython=True)
def compute_next_count(A_vals, mask):
    n = len(A_vals)
    next_count = np.zeros(n, dtype=np.int32)
    for i in range(n):
        cnt = 0
        for j in range(i+1, n):
            if mask[j] and A_vals[j] > A_vals[i]:
                cnt += 1
            else:
                break
        next_count[i] = cnt
    return next_count

# 用numba编译previous count的计算函数
@jit(nopython=True)
def compute_prev_count(A_vals, mask):
    n = len(A_vals)
    prev_count = np.zeros(n, dtype=np.int32)
    for i in range(n):
        cnt = 0
        for j in range(i-1, -1, -1):
            if mask[j] and A_vals[j] > A_vals[i]:
                cnt += 1
            else:
                break
        prev_count[i] = cnt
    return prev_count

# 转换为numpy数组传入编译后的函数
A_np = df['A'].values
mask_np = (df['A'] < df['B']).values

df['next count'] = compute_next_count(A_np, mask_np)
df['previous count'] = compute_prev_count(A_np, mask_np)

这个方案对于25k行的数据来说，几乎是瞬间就能完成计算，效率提升非常明显。

验证结果 #

运行上述代码后，得到的结果和你预期的完全一致：

A  B  next count  previous count
0  0  9           2                0
1  3  8           0                0
2  2  3           0                1
3  9  5           0                0
4  1  5           0                0
5  0  5           2                1
6  4  5           1                0
7  7  8           0                0
8  3  0           0                2
9  2  4           0                0

为什么没有完全向量化的方案？ #

这里需要说明一下：因为每一行的判断阈值是自身的A值，不同行的阈值不一样，没法用统一的向量化广播操作来处理所有行的连续计数。所以循环是最直接的方案，而通过Numpy和Numba可以把循环的效率提升到接近向量化的水平，完全能应对25k行的规模。

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