在numpy中给定索引列表添加元素的最有效方法

假设我们有一个形状为 (N, ) 的 numpy 数组 A 和一个形状为 (M, 3) 的矩阵 D，其中包含数据和另一个形状为 (M, 3) 的矩阵 I，它具有 D 中每个数据元素的对应索引。我们可以构造 A 给定的 D 和 I 以便添加重复的元素索引吗？

例子：

############# A[I] := D ###################################  
A = [0.5, 0.6]                         # Final Reduced Data Vector
D = [[0.1, 0.1 0.2], [0.2, 0.4, 0.1]]  # Data
I = [[0, 1, 0], [0, 1, 1]]             # Indices

例如：

A[0] = D[0][0] + D[0][2] + D[1][0]     # 0.5 = 0.1 + 0.2 + 0.2

由于在索引矩阵中我们有：

I[0][0] = I[0][2] = I[1][0] = 0

目标是避免循环遍历所有元素以对大 N、M (10^6-10^9) 有效。

回答

# Your example
A = [0.5, 0.6]
D = [[0.1, 0.1, 0.2], [0.2, 0.4, 0.1]]
I = [[0, 1, 0], [0, 1, 1]]

# Solution
import numpy as np    
D, I = np.array(D).flatten(), np.array(I).flatten()
print(np.bincount(I, D)) #[0.5 0.6]