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]])

有没有办法编写更多的Python代码，避免两个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毫秒。我认为这段代码可能不会成为您所做任何事情的瓶颈。

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

1 个答案: