如何向量化Pandas DataFrame行操作：基于前序行计算分位数

Pandas目前没有内置的cumquantile()方法，但我们可以通过Numba加速循环或者Pandas扩展窗口+自定义函数来高效实现需求，避免原生Python循环的低效，同时保证不出现数据泄露（仅用前序行数据计算）。

方法1：Pandas扩展窗口（代码简洁，适合小数据） #

利用expanding()窗口自动维护当前行及之前的所有数据，结合自定义函数计算分位数：

import pandas as pd
import numpy as np

np.random.seed(42)
test_df = pd.DataFrame({'column_1': np.random.random(10)})
# 定义分位点：0.1到0.9，步长0.1
quantile_points = np.arange(0.1, 1.0, 0.1)

def calculate_window_quantiles(window):
    # 对当前窗口计算指定分位数
    return np.quantile(window, quantile_points)

# 应用扩展窗口，将结果转为DataFrame并命名列
quantile_df = test_df['column_1'].expanding()\
    .apply(calculate_window_quantiles, raw=True)\
    .apply(pd.Series)
quantile_df.columns = [f'q{j}' for j in range(1, 10)]

# 合并到原DataFrame
test_df = pd.concat([test_df, quantile_df], axis=1)

说明 #

• 代码简洁易读，expanding()自动处理前序行的窗口范围
• 缺点：本质还是Python级循环，数据量大时（如n>1000）速度较慢

方法2：Numba加速循环（性能最优，适合大数据） #

用Numba将循环编译为机器码，比原生Python循环快几十到上百倍，同时内存占用低：

import pandas as pd
import numpy as np
from numba import jit

np.random.seed(42)
test_df = pd.DataFrame({'column_1': np.random.random(10)})
x = test_df['column_1'].values
quantile_points = np.arange(0.1, 1.0, 0.1)
n_rows = len(x)
n_quantiles = len(quantile_points)

@jit(nopython=True)
def cumulative_quantiles(arr, qs):
    n = len(arr)
    n_qs = len(qs)
    result = np.zeros((n, n_qs))
    # 维护动态排序的数组，每次插入新元素时保持有序
    sorted_arr = np.zeros(n)
    sorted_arr[0] = arr[0]
    # 第一行所有分位数等于自身
    result[0] = sorted_arr[0]
    
    for i in range(1, n):
        val = arr[i]
        # 找到新元素在已排序数组中的插入位置
        insert_pos = np.searchsorted(sorted_arr[:i], val)
        # 插入元素，保持数组有序
        sorted_arr[insert_pos:i+1] = np.concatenate([[val], sorted_arr[insert_pos:i]])
        # 计算当前行的所有分位数（遵循线性插值规则）
        for j in range(n_qs):
            q = qs[j]
            idx = q * i  # 窗口内有i+1个元素，索引范围0~i
            if idx.is_integer():
                result[i, j] = sorted_arr[int(idx)]
            else:
                lower_idx = int(np.floor(idx))
                upper_idx = lower_idx + 1
                result[i, j] = sorted_arr[lower_idx] + (idx - lower_idx) * (sorted_arr[upper_idx] - sorted_arr[lower_idx])
    return result

# 计算累积分位数并转为DataFrame
quantile_results = cumulative_quantiles(x, quantile_points)
quantile_df = pd.DataFrame(quantile_results, columns=[f'q{j}' for j in range(1, 10)])

# 合并到原DataFrame
test_df = pd.concat([test_df, quantile_df], axis=1)

说明 #

• Numba的@jit(nopython=True)装饰器将循环编译为机器码，大幅提升速度
• 维护有序数组避免重复排序，进一步优化性能
• 内存占用为O(n)，适合处理十万级以上的数据集

方法3：Numpy广播（仅适合极小数据） #

通过构造上三角矩阵实现纯向量化操作，但会占用O(n²)内存，仅适合n<1000的场景：

import pandas as pd
import numpy as np

np.random.seed(42)
test_df = pd.DataFrame({'column_1': np.random.random(10)})
x = test_df['column_1'].values
quantile_points = np.arange(0.1, 1.0, 0.1)
n_rows = len(x)

# 构造上三角矩阵，每行包含前i+1个元素，其余补NaN
arr = np.full((n_rows, n_rows), np.nan)
arr[np.triu_indices(n_rows)] = x.repeat(n_rows).reshape(n_rows, n_rows)[np.triu_indices(n_rows)]

# 每行排序后计算分位数
sorted_arr = np.sort(arr, axis=1)
quantile_results = np.array([np.quantile(sorted_arr[i, :i+1], quantile_points) for i in range(n_rows)])

quantile_df = pd.DataFrame(quantile_results, columns=[f'q{j}' for j in range(1, 10)])
test_df = pd.concat([test_df, quantile_df], axis=1)

说明 #

• 纯向量化操作，但内存消耗随数据量平方增长，仅适合极小数据集

性能对比 #

内容的提问来源于stack exchange，提问作者Felipe Sêrro