在 2D numpy 数组的子矩阵上高效运行

Question

我无法有效地执行以下矩阵运算。从一个方阵（2D numpy 数组）和跨越矩阵每个索引的组（字典：键是组，值是组中矩阵索引的列表）开始，我需要获得一个新的、更小的矩阵包含原始矩阵的每个子矩阵中元素的总和。子矩阵是根据组的索引定义的。因此，新矩阵也将是方形的，但以组数作为其维度。

让我们看看下面的例子：

import numpy as np

X = np.arange(49).reshape((7, 7))

d = {0: [0, 1], 1: [2, 3, 4], 2: [5, 6]}

def get_new_matrix(matrix, groups_indexes):

    groups_number = len(groups_indexes)

    new_matrix = np.zeros((groups_number, groups_number))

    for i in range(groups_number):

        for j in range(groups_number):

            new_matrix[i][j] = np.sum(matrix[groups_indexes[i]][:,groups_indexes[j]])

    return new_matrix

Z = get_new_matrix(X, d)

print(Z)

[[ 16  39  36]

 [129 216 159]

 [156 249 176]]

查看结果，例如在（第二）第 1 行和（第三）列 2 中，我们注意到结果是 159，这是：

Z[1,2]

这意味着在原始矩阵中，由行中的组 1 和列中的组 2 定义的子矩阵，即第 2、3 和 4 行以及第 5 和第 6 列，明确为：

X[[2, 3, 4]][:,[5, 6]]

并且子矩阵中所有元素的总和为19+20+26+27+33+34=159。

明确：

np.sum(X[[2, 3, 4]][:,[5, 6]])

有没有什么办法可以写出更pythonic的代码，避免两次for循环获取新矩阵，提高整体效率？我想它应该是花哨的索引、广播等，但我还没有找到更好的解决方案。

我当前的代码对于大型初始矩阵（以及潜在的大量初始组）的缩放非常糟糕，而且由于我不仅要为任意的大型初始方阵运行它，而且还要在多次迭代中运行它，我真的需要改进它。或者也许没有办法让代码变得更好，解释也会非常有用:)

Answer 1

如果您的组索引跨越整个矩阵并且是连续的，您可以将它们存储为索引而不是字典。由于每组都以下一组的开头结束，因此您只需要存储开始索引。您的当前d可以重写为d = sorted(val[0] for val in d.values())或者，如果您不依赖于字典格式，只需d = np.array([0, 2, 5])我的建议是应用np.add.reduceat两次，每个维度一次，基本上就像您在当前循环中所做的那样，但让 numpy 在内部为您管理循环：result = np.add.reduceat(np.add.reduceat(X, d, axis=0), d, axis=1)问题中输入的结果是：array([[ 16,  39,  36],       [129, 216, 159],       [156, 249, 176]])159 确实是索引 [1, 2] 处的元素。这似乎可以很好地扩展。在我不太强大的笔记本电脑上运行X = np.arange(10**6).reshape(10**3, 10**3)并d = np.arange(0, 10**3, 10)花费大约 2.27 毫秒。我不认为这段代码可能会成为你所做的任何事情的瓶颈。