如何按首元素分组原NumPy二维数组生成新二维数组

首先明确两个生成规整数组的隐含前提：相同首元素的条目在原数组中连续排列，且每个分组对应的第二列值的数量完全一致。如果这两个条件不满足，无法生成规整的纯数值二维NumPy数组（NumPy要求二维数组每行长度统一），只能存储为对象类型的变长结构。

分场景实现方法 #

场景1：满足规整数组生成条件（同组条目连续、组内元素数相等） #

优先选择纯NumPy实现，无额外依赖、运行性能更高：

import numpy as np
# 示例原数组
src_arr = np.array([
    [113, 1915], [113, 3314], [113, 53081],
    [97267, 4112], [97267, 37556], [97267, 4786]
])
# 若不能确定相同首元素的条目是连续的，先按第一列排序，规避分组错误
src_arr = src_arr[src_arr[:, 0].argsort()]
# 统计每个分组的元素数量
unique_keys, group_sizes = np.unique(src_arr[:, 0], return_counts=True)
# 校验所有组的长度一致，避免重塑报错
if not np.all(group_sizes == group_sizes[0]):
    raise ValueError("各分组长度不统一，无法生成规整二维数值数组")
# 将第二列按分组维度重塑为目标结构
res_arr = src_arr[:, 1].reshape(len(unique_keys), group_sizes[0])

如果追求快速开发、不在意Pandas依赖，也可以用更简洁的Pandas逻辑实现：

import pandas as pd
src_arr = np.array([
    [113, 1915], [113, 3314], [113, 53081],
    [97267, 4112], [97267, 37556], [97267, 4786]
])
df = pd.DataFrame(src_arr, columns=["group_key", "val"])
res_arr = df.groupby("group_key")["val"].apply(list).to_numpy()
# 仅当各列表长度一致时，可转成标准规整二维数组
res_arr = np.array(res_arr.tolist())

场景2：各分组的组内元素数量不一致 #

这种场景不需要强行生成规整二维数组，NumPy会自动将结果存为dtype=object的一维结构，每个元素对应一个组的变长值列表，是NumPy对变长分组的原生支持方式，实现逻辑和上述Pandas代码一致，去掉最后一步转规整数组的代码即可。

高频踩坑提醒：如果跳过排序步骤直接重塑，很容易把跨组的元素错误合并，导致最终分组结果和预期完全不符。无论原数组看起来是否按首元素排好，加一行排序代码的性能损耗几乎可以忽略，能规避绝大多数的结果错误。

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