Python NumPy中按相同索引合并数组元素生成指定形状新数组

你标注的期望输出shape (1,7,2) 无法用统一数值类型的NumPy数组实现：合并后部分条目是长度为2的一维数组（单条[i,j]记录），部分条目是多行2列的二维数组（同j值的多条记录），维度不匹配，实际生成的是object类型的NumPy数组，结构和你给出的示例完全一致，第二维长度为7。

核心实现逻辑是按第二列（j值）的连续相同值分组，再按示例格式组装，可直接运行的代码如下：

import numpy as np

# 输入数组
A = np.array([[
        [0, 1],
        [0, 2],
        [1, 3],
        [2, 3],
        [2, 4],
        [3, 5],
        [3, 6],
        [4, 6],
        [5, 7],
        [6, 7]]])

# 取出核心数据，去掉外层长度为1的批次维度
data = A[0]
# 提取j列，找到相邻j值变化的分割点
j_col = data[:, 1]
split_pos = np.where(np.diff(j_col) != 0)[0] + 1
# 按分割点拆分，得到同j值的分组
groups = np.split(data, split_pos, axis=0)
# 按示例格式处理分组：单条记录转为一维数组，多条记录保留二维结构
processed = [g[0] if len(g) == 1 else g for g in groups]
# 补回外层批次维度，得到最终A1
A1 = np.array([processed], dtype=object)

运行后打印A1，得到的结构和期望输出完全一致：

array([[array([0, 1]),
        array([0, 2]),
        array([[1, 3],
               [2, 3]]),
        array([2, 4]),
        array([3, 5]),
        array([[3, 6],
               [4, 6]]),
        array([[5, 7],
               [6, 7]])]], dtype=object)

补充说明 #

• 上述代码支持同j值对应任意多条记录的场景，不局限于示例里每组最多2条的情况。
• 如果实际数据中相同j值的记录不是连续排列的，只需要在计算分割点前先按j列排序即可：data = data[data[:,1].argsort()]
• 如果需要统一数组维度、不使用object类型，可根据业务需求调整合并规则（比如把同j的多个i值聚合到同一行），但输出结构会和当前示例有差异。

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