如何用向量化方式替代循环，基于分箱结果归类变量至binned_list？

优化后的代码 #

import numpy as np
N_bins = 3
Bins = np.linspace(0, 10, N_bins+1)

binned_list = [[] for b in range(N_bins)]

N_elements = 5

np.random.seed(1)

for k in range(3):
    x = np.random.random((N_elements,2))/np.sqrt(2)*10
    mod_x = np.sqrt(x[:,0]**2 + x[:,1]**2)
    dig_x = np.digitize(mod_x, bins=Bins) - 1
    
    # 用向量化筛选替代逐个元素循环
    for bin_idx in range(N_bins):
        # 生成当前分箱的布尔掩码
        mask = dig_x == bin_idx
        # 批量添加符合条件的元素到对应分箱
        binned_list[bin_idx].extend(x[mask])

原理与优势 #

原代码内层循环逐个遍历x元素并执行append，效率较低。优化方案利用numpy的布尔索引实现向量化处理：

1. 针对每个分箱索引，生成布尔掩码mask，标记x中属于当前分箱的所有元素
2. 通过x[mask]批量提取该分箱的元素，再用extend一次性添加到binned_list对应位置

这种方式把内层循环次数从元素数量（每次5次）缩减到分箱数量（固定3次），且numpy的向量化操作由底层C实现，比Python级循环高效得多，当N_elements规模较大时，效率提升会非常显著。同时该方法完全保留了原代码的元素添加顺序，输出结果与原代码一致。

为什么不能直接用binned_list[dig_x].append(x) #

• binned_list是列表的列表，不支持数组索引（dig_x是numpy数组），无法通过数组批量定位分箱
• 若改用numpy数组存储binned_list，由于每个分箱的元素数量不固定，无法创建形状统一的数组，因此列表的列表仍是最适配的存储结构

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