本文介绍一种纯 NumPy 向量化方法，无需 Python 循环，即可根据标签数组 y 的唯一值将特征数组 x 拆分为多个子数组，并以字典形式组织结果。
本文介绍一种纯 numpy 向量化方法，无需 python 循环，即可根据标签数组 `y` 的唯一值将特征数组 `x` 拆分为多个子数组，并以字典形式组织结果。
在机器学习与数据预处理中，常需按类别标签对样本矩阵进行分组——例如将训练数据按类别索引切分为若干子集用于独立分析或交叉验证。若使用 for 循环或列表推导式逐个筛选，不仅代码冗长，且在大数据量下性能较差。幸运的是，NumPy 提供了一套高效、可组合的向量化工具，能以简洁一行（逻辑上）完成该任务。
以下是最推荐的实现方式：
import numpy as np x = np.array([[2,3,5,6], [1,2,4,3], [1,5,6,4], [2,8,9,5]]) y = np.array([1,0,1,2]) # 核心一行：按 y 值分组 x n, cnt = np.unique(y, return_counts=True) out = dict(zip(n, np.array_split(x[np.argsort(y)], np.cumsum(cnt[:-1])))) print(out)
输出为：
{
0: array([[1, 2, 4, 3]]),
1: array([[2, 3, 5, 6],
[1, 5, 6, 4]]),
2: array([[2, 8, 9, 5]])
}✅ 原理简析（三步完成）：
- 排序对齐：np.argsort(y) 返回 y 升序排列的索引，x[np.argsort(y)] 将 x 行按 y 值从小到大重排；
- 定位分割点：np.unique(y, return_counts=True) 获取各标签出现次数 cnt，np.cumsum(cnt[:-1]) 计算前缀和（即每个新类别的起始位置）；
- 无损切分：np.array_split(..., indices_or_sections) 按指定位置切分重排后的 x，再用 zip(n, ...) 关联标签与对应子数组。
⚠️ 注意事项：
- np.array_split 在切分点超出数组长度时仍安全（自动忽略多余分割），比 np.split 更鲁棒；
- 若 y 中存在非连续整数（如 [10, 5, 10, 15]），结果键仍为原始值（{5: ..., 10: ..., 15: ...}），无需映射；
- 所有操作均为向量化，时间复杂度近似 $O(N \log N)$（主要来自 argsort），远优于 $O(N \cdot K)$ 的循环方案。
? 进阶建议：
若需返回列表而非字典（如保持顺序），可改用 list(np.array_split(...)) 并配合 np.unique(y, return_index=True) 获取原始顺序；若后续需频繁索引某类，建议封装为函数并添加类型提示，提升可维护性。
此方法兼顾简洁性、性能与可读性，是 NumPy 高级索引与分组操作的典型范例。
