Pandas百万级数据：满足条件时高效提取可变区间最大height

针对你百万级DataFrame的需求，纯Python循环效率极低，下面提供两种高效实现方案：

方案一：Numba JIT编译加速循环 #

Numba能将Python循环编译为机器码，彻底消除解释型语言的性能开销，是这类场景下最直接的提速方案。

import numba
import pandas as pd
import numpy as np

@numba.jit(nopython=True)
def compute_highest(height_arr, cond1_arr, distance_arr):
    n = len(height_arr)
    highest = np.zeros(n, dtype=np.float64)
    for i in range(n):
        if cond1_arr[i]:
            end_idx = i + distance_arr[i]
            # 防止索引超出数组边界
            end_idx = min(end_idx, n - 1)
            # 确保区间至少包含一个元素
            if end_idx > i:
                highest[i] = np.max(height_arr[i+1:end_idx+1])
            else:
                highest[i] = 0.0
        else:
            highest[i] = 0.0
    return highest

# 直接操作numpy数组，避免Pandas额外开销
df['highest'] = compute_highest(df['height'].values, df['cond1'].values, df['distance'].values)

关键说明 #

• nopython=True 是核心：让Numba生成完全脱离Python解释器的机器码，性能提升最显著
• 直接处理numpy数组而非Pandas列，减少中间转换的开销
• 加入边界判断，避免索引越界导致的错误

方案二：稀疏表实现RMQ（范围最大值查询） #

如果不想依赖Numba，可以用稀疏表做预处理，实现O(1)时间的区间最大值查询，整体复杂度为O(nlogn)预处理 + O(n)查询，适合大规模数据。

import numpy as np

def build_sparse_table(arr):
    n = len(arr)
    k = np.floor(np.log2(n)).astype(int)
    # 构建稀疏表，st[j][i]表示从i开始，长度为2^j的区间的最大值
    st = np.zeros((k+1, n), dtype=arr.dtype)
    st[0] = arr
    for j in range(1, k+1):
        for i in range(n - (1 << j) + 1):
            st[j][i] = max(st[j-1][i], st[j-1][i + (1 << (j-1))])
    return st, k

def query_max(st, k, l, r):
    # 查询[l, r]区间的最大值
    length = r - l + 1
    j = np.floor(np.log2(length)).astype(int)
    return max(st[j][l], st[j][r - (1 << j) + 1])

# 预处理height数组的稀疏表
height_arr = df['height'].values
st, k = build_sparse_table(height_arr)

# 计算每个行的有效右边界，防止越界
end_idx = df.index + df['distance']
end_idx = end_idx.clip(upper=len(df)-1)

# 初始化highest列为0
df['highest'] = 0.0

# 筛选需要查询的行（cond1为True且区间有效）
mask = df['cond1'] & (end_idx > df.index)

# 批量查询区间最大值
df.loc[mask, 'highest'] = [
    query_max(st, k, i+1, end) 
    for i, end in zip(df.index[mask], end_idx[mask])
]

关键说明 #

• 稀疏表预处理只需要一次，后续查询都是O(1)
• 用mask筛选需要处理的行，避免无意义的计算
• 适合需要多次进行区间查询的场景

性能参考 #

• 纯Python循环：百万级数据可能需要数十分钟
• Numba方案：通常几秒到十几秒完成
• 稀疏表方案：预处理+查询耗时几秒，具体取决于数据规模

注意事项 #

• 确保distance列是整数类型，避免索引计算出错
• 当distance为0或负数时，区间无效，直接赋值为0

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