如何高效执行基于索引数组的numpy元素堆叠操作？

嘿，我来帮你搞定这个根据索引处理numpy数组元素的需求！首先先纠正个小细节：你给出的数组a定义少了一层括号，正确写法应该是a = numpy.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]])，不然会抛出维度不匹配的错误哦。

简单来说，我们要对数组a的每一行，按照array_index对应位置的索引值，把a的元素"归位"到目标位置——如果多个元素指向同一个索引，就把它们聚合起来（默认是求和，也可以改成堆叠成列表）。

1. 基础求和聚合（统一或可变行长度） #

先导入numpy，再定义正确的数组：

import numpy as np

# 修正后的数组a
a = np.array([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]])
# 给定的索引数组
array_index = np.array([[0, 0, 1, 2], [0, 1, 2, 2], [0, 1, 1, 3]])

如果想要每行长度和目标索引匹配（无多余0），可以这么写：

result = []
for row_idx in range(a.shape[0]):
    # 计算当前行需要的结果长度（最大索引+1）
    target_length = array_index[row_idx].max() + 1
    # 初始化当前行的结果数组
    row_result = np.zeros(target_length, dtype=a.dtype)
    # 用np.add.at实现索引位置的累加，这是处理这类聚合的关键方法
    np.add.at(row_result, array_index[row_idx], a[row_idx])
    result.append(row_result)

# 转成numpy数组（因为每行长度不同，会用object类型存储）
result = np.array(result, dtype=object)

运行后得到的结果是：

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

解释下逻辑：

• 第一行：a[0]的1、2都指向索引0，累加得3；3指向索引1，保留3；4指向索引2，保留4
• 第二行：a[1]的1指向索引0，保留1；2指向索引1，保留2；3和4都指向索引2，累加得7
• 第三行：a[2]的1指向索引0，保留1；2、3都指向索引1，累加得5；4指向索引3，保留4

如果想要统一行长度（补0填充），只需要调整初始化逻辑：

# 先计算所有行中最大的目标长度
max_length = array_index.max() + 1
result = np.zeros((a.shape[0], max_length), dtype=a.dtype)

for row_idx in range(a.shape[0]):
    np.add.at(result[row_idx], array_index[row_idx], a[row_idx])

得到的结果是：

array([[3, 3, 4, 0],
       [1, 2, 7, 0],
       [1, 5, 0, 4]])

2. 元素堆叠成列表（而非求和） #

如果你的需求不是求和，而是把同一目标位置的元素都存成列表，可以用字典来分组：

result = []
for row_idx in range(a.shape[0]):
    idx_group = {}
    # 遍历当前行的元素和对应索引
    for val, idx in zip(a[row_idx], array_index[row_idx]):
        if idx not in idx_group:
            idx_group[idx] = []
        idx_group[idx].append(val)
    # 按索引从小到大排序，保证顺序正确
    sorted_row = [idx_group[k] for k in sorted(idx_group.keys())]
    result.append(sorted_row)

运行后得到：

[[[1, 2], [3], [4]], [[1], [2], [3, 4]], [[1], [2, 3], [4]]]

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