python - 在 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 毫秒。我不认为这段代码可能成为您所做的任何事情的瓶颈。