如何按照另一数组分箱后的子数组长度对numpy数组进行切片

实现方案 #

方案1：NumPy 原生分割（最适配当前场景，简洁高效） #

直接用np.split配合前缀和计算分割点，即可批量完成y的分组：

# 计算各分箱长度的前缀和作为分割位置
split_points = np.cumsum([len(sub_arr) for sub_arr in x_bin])[:-1]
# 分割得到y的分组列表，y_bins[i] 对应你需要的yi
y_bins = np.split(y, split_points)

方案2：分箱阶段同步获取y分组（更通用的最优实践） #

你当前是先给x分箱再切y，其实可以在分箱时直接同步拿到对应y的分组，避免后续二次处理，同时也适配x无序的场景：

bins = np.arange(0, 1, 0.1)
bin_ids = np.digitize(x, bins)
x_bin = []
y_bins = []
for i in range(1, len(bins)):
    mask = bin_ids == i
    x_bin.append(x[mask])
    y_bins.append(y[mask])

注意：你当前手动按长度切片连续y的逻辑只适用于x是有序数组、分箱结果连续的场景，如果x本身是乱序的，按长度切连续y会导致分组和x的分箱不对应，方案2的写法不存在这个问题，通用性更强。

方案3：纯Python通用实现（不依赖NumPy） #

如果需要纯Python写法，可以用迭代器切片实现：

from itertools import islice
y_iter = iter(y)
y_bins = [list(islice(y_iter, len(sub_arr))) for sub_arr in x_bin]

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