如何根据分组但未排序的NumPy数组y的值高效拆分NumPy数组x

针对你提出的需求——根据未排序的标签数组y拆分特征数组x，用NumPy实现的最高效方式肯定是利用向量化操作避免Python循环，毕竟循环在处理大数据时会拖慢速度。下面给你两种实用方案，按需选择：

方案一：布尔索引直接分组（直观易读） #

这种方法逻辑最直白，适合标签类别不多的场景：

1. 先把你的列表转成NumPy数组（充分利用NumPy的底层优化优势）：

import numpy as np

x = np.array([[1, 2, 8], [2, 9, 1], [3, 8, 9], [4, 3, 5], [5, 2, 3], [6, 4, 7], [7, 2, 3], [8, 2, 2], [9, 5, 3], [10, 2, 3], [11, 2, 4]])
y = np.array([0, 0, 1, 0, 1, 1, 2, 2, 2, 0, 0])

2. 获取y中所有唯一标签（默认按升序排列）：

unique_labels = np.unique(y)

3. 遍历标签，用布尔索引提取对应行：

# 用字典存储分组结果，键是标签值，值是对应的子数组
groups = {}
for label in unique_labels:
    groups[f"z_{label}"] = x[y == label]

# 输出查看结果
print(groups["z_0"])
print(groups["z_1"])
print(groups["z_2"])

这个方案的优势是代码清晰易懂，而且布尔索引是NumPy的底层优化操作，速度远快于Python原生循环，还能保留原数组中元素的相对顺序。

方案二：排序后拆分（大数据量更高效） #

如果你的数组规模很大，标签类别也比较多，这种方法会更高效——通过一次排序+拆分完成操作，减少多次索引的开销：

1. 先对标签数组y排序，得到排序后的索引（NumPy的argsort默认是稳定排序，相同标签的元素会保留原相对顺序）：

sorted_indices = np.argsort(y)
sorted_x = x[sorted_indices]
sorted_y = y[sorted_indices]

2. 找到标签变化的位置，作为拆分点：

# np.diff计算相邻元素差值，不为0的位置就是标签切换的地方，加1是为了得到拆分的起始索引
split_positions = np.where(np.diff(sorted_y))[0] + 1

3. 用np.split一次性拆分数组：

z_0, z_1, z_2 = np.split(sorted_x, split_positions)

这样得到的z_0、z_1、z_2完全符合你的需求，而且在大数据量下，这种单次批量操作的效率会比多次布尔索引更高。

为什么这两种方法高效？ #

NumPy的向量化操作是基于C语言实现的，避免了Python循环的解释器开销，处理大规模数据时速度会比纯Python代码快几个数量级。这两种方案都没有用到Python原生循环，完全依赖NumPy的内置函数，是处理这类分组问题的最优选择。

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