Python中如何高效实现自定义区间分箱并构建数组？

原代码的核心问题是嵌套循环重复遍历全量数据，时间复杂度为O(n*m)（n为数据量，m为分箱数），16万条数据+450个分箱会产生7200万次判断，效率自然低下。利用pandas的内置向量化操作可以将时间复杂度降到O(n log m)甚至O(n)，大幅提升效率。

以下是两种优化方案：

方案一：利用pd.cut实现分箱分组 #

pd.cut是pandas专门用于区间分箱的函数，配合groupby可以快速完成分箱收集：

import pandas as pd

# 读取数据并命名列（方便操作）
dfrawdata = pd.read_excel(r"data.xlsx", names=['x', 'y'])
dfbins = pd.read_excel(r"bins.xlsx", usecols=[1, 3], names=['bin_left', 'bin_right'])

# 整理分箱区间：将左右边界转为区间数组
bins_intervals = pd.IntervalIndex.from_arrays(dfbins['bin_left'], dfbins['bin_right'], closed='neither')

# 为每条数据匹配对应的分箱标签
dfrawdata['bin_label'] = pd.cut(dfrawdata['x'], bins=bins_intervals)

# 按分箱分组，收集y值（结果为字典，键是区间，值是对应y列表）
bin_results = dfrawdata.groupby('bin_label')['y'].apply(list).to_dict()

# 如果需要和原分箱索引对应，可以转换为列表形式
emptyarray = [bin_results.get(interval, []) for interval in bins_intervals]

方案二：利用pd.merge_asof实现高效匹配（适合已排序数据） #

由于你的x数据是按递增排序的，merge_asof可以在O(n + m)时间复杂度内完成分箱匹配，效率比pd.cut更高：

import pandas as pd

# 读取数据并命名列
dfrawdata = pd.read_excel(r"data.xlsx", names=['x', 'y'])
dfbins = pd.read_excel(r"bins.xlsx", usecols=[1, 3], names=['bin_left', 'bin_right'])

# 为分箱添加索引，后续用于匹配
dfbins['bin_index'] = dfbins.index

# 使用merge_asof匹配：找到每个x对应的第一个bin_right大于x的分箱
# 注意dfrawdata和dfbins都需要按x/bin_left排序（你的数据已经满足，若不确定可添加sort操作）
merged = pd.merge_asof(
    dfrawdata,
    dfbins,
    left_on='x',
    right_on='bin_left',
    direction='forward'
)

# 过滤出x在分箱区间内的数据（排除x大于所有bin_right的情况）
merged = merged[merged['x'] < merged['bin_right']]

# 按分箱索引分组，收集y值
bin_groups = merged.groupby('bin_index')['y'].apply(list)

# 构建结果列表，确保每个分箱位置都有对应的列表（空列表表示无匹配数据）
emptyarray = [bin_groups.get(idx, []) for idx in dfbins.index]

关键优化点说明 #

• 替换嵌套循环为pandas向量化操作：向量化操作由C语言底层实现，比Python循环快几个数量级
• 利用数据排序特性：merge_asof针对有序数据做近似匹配，避免了全量遍历
• 避免重复数据读取/判断：原代码每次遍历分箱都要重新读一遍全量数据，优化后只需要处理一次

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