如何优化Python迭代？高效生成价格区间统计DataFrame

背景 #

pivots是由以下代码生成的numpy数组列表：

pivots = []
for i in range(1, dimension):
    price_range = np.linspace(a, b, dimension//i).reshape(-1, 1)
    pdf = np.exp(kde.score_samples(price_range))
    peaks = find_peaks(pdf)[0]
    pivot_levels = price_range[peaks]
    pivots.append(pivot_levels)

需求 #

生成包含两列的Pandas DataFrame：

• 第一列：步长为200的价格区间，格式如30,000 - 30,200（下限值 - 上限值）
• 第二列：统计每个numpy数组中落在对应价格区间内的数值数量（原代码逻辑为每个pivot数组的每个区间对应一行记录）

原代码性能瓶颈 #

原代码存在两个核心性能问题：

1. 双重循环+逐行追加DataFrame：df.append()每次都会创建新的DataFrame实例，时间复杂度为O(n²)，循环次数越多性能越差
2. 逐区间判断计数：对每个pivot数组的每个区间单独做布尔判断求和，没有利用numpy的向量化统计能力

原代码如下：

# Create an empty DataFrame
df = pd.DataFrame(columns=['Price Range', 'Count'])

# Define the price range step and format
step = 200
price_format = "{:,.0f} - {:,.0f}"

# Iterate over the pivots list
for i, pivot_levels in enumerate(pivots):
    # Iterate over the price ranges
    for lower_range in np.arange(a, b, step):
        upper_range = lower_range + step
        # Format the price range string
        price_range_str = price_format.format(lower_range, upper_range)
        # Count the number of values in the pivot_levels that fall within the current range
        count = np.sum((lower_range <= pivot_levels) & (pivot_levels < upper_range))
        # Append a new row to the DataFrame
        df = df.append({'Price Range': price_range_str, 'Count': count}, ignore_index=True)

优化方案 #

核心优化思路 #

1. 预先生成所有区间的边界和格式化字符串，避免重复计算
2. 用np.histogram向量化统计每个pivot数组在所有区间的数量，替代逐区间判断
3. 批量生成每个pivot对应的DataFrame片段，最后一次性合并，替代逐行追加

优化后的代码 #

import numpy as np
import pandas as pd

step = 200
price_format = "{:,.0f} - {:,.0f}"

# 1. 预先生成所有价格区间的边界和格式化字符串
bins = np.arange(a, b + step, step)  # 包含上限，确保最后一个区间被完整覆盖
price_ranges = [price_format.format(bins[i], bins[i+1]) for i in range(len(bins)-1)]

# 2. 批量处理每个pivot数组，收集所有结果片段
result_frames = []
for pivot_levels in pivots:
    # 用histogram一次性统计当前pivot在所有区间的数量
    counts, _ = np.histogram(pivot_levels.flatten(), bins=bins)
    # 生成当前pivot对应的DataFrame片段
    temp_df = pd.DataFrame({
        'Price Range': price_ranges,
        'Count': counts
    })
    result_frames.append(temp_df)

# 3. 一次性合并所有片段，高效生成最终DataFrame
df = pd.concat(result_frames, ignore_index=True)

性能提升说明 #

• 向量化统计：np.histogram底层用C实现，比Python循环逐区间求和快10~100倍
• 批量合并：pd.concat是批量操作，比多次df.append效率提升几个数量级，尤其是当pivots数量和区间数较多时

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