python - 在 Python 中按索引对矩阵元素求和

Question

我有两个矩阵(相同的行和列):一个具有浮点值，在另一个矩阵中按索引分组。因此，我想要一个字典或一个列表，其中包含每个索引的元素总和。 索引始终从 0 开始。

A = np.array([[0.52,0.25,-0.45,0.13],[-0.14,-0.41,0.31,-0.41]])
B = np.array([[1,3,1,2],[3,0,2,2]])

RESULT = {0: -0.41, 1: 0.07, 2: 0.03, 3: 0.11}

我找到了这个解决方案，但我正在寻找更快的解决方案。 我正在处理具有 784 x 300 个单元格的矩阵，此算法需要大约 28 毫秒才能完成。

import numpy as np

def matrix_sum_by_indices(indices,matrix):
    a = np.hstack(indices)
    b = np.hstack(matrix)
    sidx = a.argsort()
    split_idx = np.flatnonzero(np.diff(a[sidx])>0)+1
    out = np.split(b[sidx], split_idx)
    return [sum(x) for x in out]

如果你能帮我找到一个更好更简单的解决方案来解决这个问题，我将不胜感激!

编辑:我犯了一个错误，完成时间在 300*10 矩阵中约为 8 毫秒，但在 784x300 中约为 28 毫秒。

EDIT2:我的 A 元素是 float64，所以 bincount 给我 ValueError。

Answer 1

您可以在此处使用 bincount:

a = np.array([[0.52,0.25,-0.45,0.13],[-0.14,-0.41,0.31,-0.41]])
b = np.array([[1,3,1,2],[3,0,2,2]])

N = b.max() + 1
id = b + (N*np.arange(b.shape[0]))[:, None] # since you can't apply bincount to a 2D array
np.sum(np.bincount(id.ravel(), a.ravel()).reshape(a.shape[0], -1), axis=0)

输出:

array([-0.41,  0.07,  0.03,  0.11])

作为函数:

def using_bincount(indices, matrx):
    N = indices.max() + 1
    id = indices + (N*np.arange(indices.shape[0]))[:, None] # since you can't apply bincount to a 2D array
    return np.sum(np.bincount(id.ravel(), matrx.ravel()).reshape(matrx.shape[0], -1), axis=0)

此示例的时间:

In [5]: %timeit using_bincount(b, a)
31.1 µs ± 1.74 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [6]: %timeit matrix_sum_by_indices(b, a)
61.3 µs ± 2.62 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [88]: %timeit scipy.ndimage.sum(a, b, index=[0,1,2,3])
54 µs ± 218 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

(scipy.ndimage.sum 在更大的样本上应该更快)