Python中按索引分组求矩阵元素之和

Question

Python中按索引分组求矩阵元素之和

6

我有两个矩阵（行数和列数相同）：一个包含浮点值，这些值按照另一个矩阵中的索引进行分组。因此，我想要一个字典或列表，其中包含每个索引的元素总和。

索引始终从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的矩阵中完成的时间是~8ms，但在784x300的矩阵中完成的时间是~28ms。编辑2：我的"A"元素是float64类型，所以bincount会给我一个ValueError。

- Mortafix

8毫秒？我会说那相当快了。你需要多快？ - Bayko

@Bayko 我正在寻找类似于 ms 下的东西。因为这个过程在一个大约 6000*100 的循环中。 - Mortafix

你可能正在寻找某种形式的 **bincount**。 - user3483203

3个回答

1

numpy_indexed包提供了高效且简单的解决方案来解决这个问题（免责声明：我是它的作者）：

import numpy_indexed as npi
keys, values = npi.group_by(B.flatten()).sum(A.flatten())

- Eelco Hoogendoorn

我该如何在Jupyter Notebook中使用它？ - Mortafix

有一个conda-forge软件包和一个pip安装程序可用；无论使用笔记本还是不使用，都不会有区别。 - Eelco Hoogendoorn

1

下面的解决方案依赖于scipy.ndimage.sum，速度高度优化：

import numpy as np
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]])
import scipy.ndimage
print(scipy.ndimage.sum(A, B, index=[0,1,2,3]))

你可能需要做一些工作，才能使index参数完全符合你的要求。它是你想在结果中获取的索引列表。以下内容可能是一个很好的起点：

print(scipy.ndimage.sum(A,B, index=np.unique(B)))

但是如果您预先知道所有索引的列表，将其硬编码在此处会更有效。

- Thomas Baruchel

我该如何在float64中使用它？ - Mortafix

你可以在数组 A 中拥有任何你想要的东西；这是你所要求的吗？ - Thomas Baruchel

没事了，它可以使用float64。它比我的快（约10毫秒），但如果可能的话，我正在寻找更快的。 - Mortafix

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- user3483203 · Accepted Answer

您可以在这里使用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 应该更快。