python - 根据第二个数组在一个数组中查找项目

Question

我有两个数组 A 和 B:

A=array([[ 5.,  5.,  5.],
         [ 8.,  9.,  9.]])
B=array([[ 1.,  1.,  2.],
         [ 3.,  2.,  1.]])

B 中任何有“1”的地方我想对 A 中相同的行和列位置求和。

例如，对于这个问题，答案是 5+5+9=10

我希望它持续 2,3....n(B 中的所有唯一值)

所以对于 2 的...它会是 9+5=14 而对于 3 的它会是 8

我通过使用找到了唯一值:

numpy.unique(B)

我意识到这 make 需要多个步骤，但我无法真正使用索引矩阵来对另一个矩阵中的这些位置求和。

Answer 1

对于每个唯一值x，你可以做

A[B == x].sum()

例子:

>>> A[B == 1.0].sum()
19.0