如何在Pandas DataFrame中高效统计前N个单元格内符合范围的数量

问题描述 #

给定如下结构的Pandas DataFrame：

col1
0     10
1      5
2      8
3     12
4     13
5      6
6      9
7     11
8     10
9      3
10    21
11    18
12    14
13    16
14    30
15    45
16    31
17    40
18    38

首先为col1列每个单元格计算数值范围：

df['df_min'] = df.col1 - df.col1 * 0.2  # 等价于 df.col1 * 0.8
df['df_max'] = df.col1 + df.col1 * 0.2  # 等价于 df.col1 * 1.2

需求是为每个单元格统计其之前3个单元格中，col1值落在该单元格对应[df_min, df_max]范围内的数量。当前使用循环实现的方案在大数据量下性能极差，需要无需循环的高效解决方案。目标结果如下：

col1  df_min  df_max  counter
0     10     8.0    12.0       -1
1      5     4.0     6.0       -1
2      8     6.4     9.6       -1
3     12     9.6    14.4        1
4     13    10.4    15.6        1
5      6     4.8     7.2        0
6      9     7.2    10.8        0
7     11     8.8    13.2        2
8     10     8.0    12.0        2
9      3     2.4     3.6        0
10    21    16.8    25.2        0
11    18    14.4    21.6        1
12    14    11.2    16.8        0
13    16    12.8    19.2        2
14    30    24.0    36.0        0
15    45    36.0    54.0        0
16    31    24.8    37.2        1
17    40    32.0    48.0        1
18    38    30.4    45.6        3

原循环实现代码 #

df = pd.DataFrame({"col1":[10, 5, 8, 12, 13, 6, 9, 11, 10, 3, 21, 18, 14, 16, 30, 45, 31, 40, 38]})

back = 3 # 需要回溯的单元格数量

df['df_min'] = df.col1 - df.col1 * 0.2
df['df_max'] = df.col1 + df.col1 * 0.2

l = []
for window in df.col1.rolling(window=back+1, center=False, closed='right'):
    if window.empty:
        pass
    else:
        a = window.iloc[-1]
        range_min = a - a * 0.2
        range_max = a + a * 0.2
        c = 0
        if len(window) == back+1:
            for b in window:
                if (b >= range_min and b <= range_max):
                    c += 1
        c = c-1 # 减去窗口中包含的当前值（必然符合条件）
        l.append(c)
df1 = pd.DataFrame(l, columns=['counter'])

df = df.join(df1)

print(df)

高效向量化解决方案 #

利用Numpy的滑动窗口和广播操作，完全避免Python层面的循环，大幅提升性能：

import pandas as pd
import numpy as np

# 初始化数据
df = pd.DataFrame({"col1":[10, 5, 8, 12, 13, 6, 9, 11, 10, 3, 21, 18, 14, 16, 30, 45, 31, 40, 38]})
back = 3

# 计算范围值
df['df_min'] = df['col1'] * 0.8
df['df_max'] = df['col1'] * 1.2

# 初始化counter列，前back行没有足够前置元素，设为-1
df['counter'] = -1

# 生成滑动窗口矩阵：每一行对应当前行之前的back个元素
# 矩阵形状为 (len(df)-back, back)
window_matrix = np.lib.stride_tricks.sliding_window_view(df['col1'].values, window_shape=back)

# 提取需要计算的行对应的min和max值
target_mins = df['df_min'].values[back:]
target_maxs = df['df_max'].values[back:]

# 广播比较：每个窗口元素与对应行的min、max比较，生成布尔矩阵
mask = (window_matrix >= target_mins[:, np.newaxis]) & (window_matrix <= target_maxs[:, np.newaxis])

# 统计每行符合条件的元素数量
counts = mask.sum(axis=1)

# 将结果赋值给counter列
df.loc[back:, 'counter'] = counts

print(df)

方案优势 #

• 无Python循环：所有核心计算都在Numpy的C级运算中完成，避免了循环带来的性能损耗。
• 时间复杂度O(n)：滑动窗口生成和广播比较的时间复杂度均为线性，处理百万级甚至更大数据集时性能远超循环方案。
• 代码简洁：逻辑清晰，易于维护和扩展（只需修改back参数即可调整回溯数量）。

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