如何使用索引数组将Numpy数组分组为多个子数组？

这问题我太懂了！要实现按数组b的取值对数组a进行分组，用NumPy就能高效搞定，不管数据规模多大都适用。给你几个实用的实现方式：

方法1：简洁字典推导式（快速实现） #

直接利用布尔索引筛选元素，一行代码就能构建目标字典：

import numpy as np

a = np.array([18, 56, 32, 75, 55, 55])
b = np.array([0, 2, 3, 2, 2, 2])

# 构建分组字典c
c = {val: a[b == val] for val in np.unique(b)}

# 验证结果
print(c[0])   # 输出: array([18])
print(c[2])   # 输出: array([56, 75, 55, 55])
print(c[3])   # 输出: array([32])

这种写法直观易懂，适合大多数日常场景。唯一要注意的是，np.unique(b)会返回排序后的唯一值，但因为我们用值作为字典的键，顺序不影响最终的分组结果。

方法2：更高效的批量索引（适合大规模数据） #

如果你的实际数据量很大，重复执行b == val会多次遍历数组，效率稍低。可以用np.unique的return_inverse参数一次性获取分组索引，再批量提取元素：

import numpy as np

a = np.array([18, 56, 32, 75, 55, 55])
b = np.array([0, 2, 3, 2, 2, 2])

# 获取唯一值和每个元素对应的分组索引
unique_vals, group_indices = np.unique(b, return_inverse=True)

c = {}
for val in unique_vals:
    # 筛选出当前分组的所有元素
    c[val] = a[group_indices == np.where(unique_vals == val)[0][0]]

# 验证结果和之前一致
print(c[0])   # array([18])
print(c[2])   # array([56, 75, 55, 55])
print(c[3])   # array([32])

这种方法只需要遍历一次数组获取分组索引，后续的筛选都是基于索引的快速操作，性能更优。

额外技巧：保持分组顺序和首次出现一致 #

如果需要字典的键顺序和b中值首次出现的顺序一致（比如示例里的0→2→3），可以用return_index参数来实现：

unique_vals, first_indices = np.unique(b, return_index=True)
# 按首次出现的顺序排序唯一值
sorted_vals = unique_vals[np.argsort(first_indices)]

c = {val: a[b == val] for val in sorted_vals}

这样字典c的键顺序就和你示例里的完全匹配啦。

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