优化测试来自多个 NumPy 数组的所有行组合

Question

我有三个 NumPy 数组 整数，列数相同，行数任意。我对第一个的一行加上第二个的一行给出第三个的行的所有实例感兴趣（[3, 1, 4] + [1, 5, 9] = [4, 6, 13 ])。

这是一个伪代码：

for i, j in rows(array1), rows(array2):
    if i + j is in rows(array3):
        somehow store the rows this occured at (eg. (1,2,5) if 1st row of 
        array1 + 2nd row of array2 give 5th row of array3)

我需要为非常大的矩阵运行这个，所以我有两个问题：

(1) 我可以使用嵌套循环编写上述内容，但是否有更快的方法，可能是 list comprehensions 或 itertools？

(2) 什么是最快/最节省内存的存储方式三元组？稍后我将需要创建一个热图，使用两个作为坐标，第一个作为相应的值，例如。在伪代码示例中，点 (2,5) 的值为 1。

非常感谢任何提示 - 我知道这听起来很简单，但它需要快速运行，而且我对优化的经验很少。

编辑：我的丑陋代码是在 cmets 中请求的

import numpy as np

#random arrays
A = np.array([[-1,0],[0,-1],[4,1], [-1,2]])
B = np.array([[1,2],[0,3],[3,1]])
C = np.array([[0,2],[2,3]])

#triples stored as numbers with 2 coordinates in a otherwise-zero matrix
output_matrix = np.zeros((B.shape[0], C.shape[0]), dtype = int)
for i in range(A.shape[0]):
    for j in range(B.shape[0]):
        for k in range(C.shape[0]):
            if np.array_equal((A[i,] + B[j,]), C[k,]):
                output_matrix[j, k] = i+1

print(output_matrix) 

Answer 1

我们可以利用broadcasting 以矢量化的方式执行所有这些求和和比较，然后在其上使用np.where 来获取与匹配的索引对应的索引，最后索引并分配 -

output_matrix = np.zeros((B.shape[0], C.shape[0]), dtype = int)

mask = ((A[:,None,None,:] + B[None,:,None,:]) == C).all(-1)

I,J,K = np.where(mask)
output_matrix[J,K] = I+1

Answer 2

(1) 改进

您可以将集合用于第三个矩阵中的最终结果，因为a + b = c 必须保持相同。这已经用恒定时间查找替换了一个嵌套循环。我将在下面向您展示如何执行此操作的示例，但我们首先应该介绍一些符号。

对于基于集合的工作方法，我们需要一个可散列类型。因此，列表将不起作用，但 tuple 会：它是一个有序的、不可变的结构。但是有一个问题：元组加法定义为追加，即

(0, 1) + (1, 0) = (0, 1, 1, 0).

这不适用于我们的用例：我们需要逐元素添加。因此，我们将内置元组子类化如下，

class AdditionTuple(tuple):

    def __add__(self, other):
        """
        Element-wise addition.
        """
        if len(self) != len(other):
            raise ValueError("Undefined behaviour!")

        return AdditionTuple(self[idx] + other[idx]
                             for idx in range(len(self)))

我们覆盖__add__ 的默认行为。现在我们有了适合我们问题的数据类型，让我们准备数据。

你给我们，

A = [[-1, 0], [0, -1], [4, 1], [-1, 2]]
B = [[1, 2], [0, 3], [3, 1]]
C = [[0, 2], [2, 3]]

合作。我说，

from types import SimpleNamespace

A = [AdditionTuple(item) for item in A]
B = [AdditionTuple(item) for item in B]
C = {tuple(item): SimpleNamespace(idx=idx, values=[])
     for idx, item in enumerate(C)}

也就是说，我们修改 A 和 B 以使用我们的新数据类型，并将 C 转换为支持（摊销）O(1) 查找时间的字典。

我们现在可以执行以下操作，完全消除一个循环，

from itertools import product

for a, b in product(enumerate(A), enumerate(B)):
    idx_a, a_i = a
    idx_b, b_j = b

    if a_i + b_j in C:  # a_i + b_j == c_k, identically
        C[a_i + b_j].values.append((idx_a, idx_b))

那么，

>>>print(C)
{(2, 3): namespace(idx=1, values=[(3, 2)]), (0, 2): namespace(idx=0, values=[(0, 0), (1, 1)])}

对于C 中的每个值，您将获得该值的索引（如idx），以及(idx_a, idx_b) 的元组列表，其中A 和B 的元素总和为该值在idx 在C。

让我们简单分析一下这个算法的复杂性。如上所述重新定义列表A、B 和C 在列表长度上是线性的。遍历A 和B 当然是在O(|A| * |B|) 中，并且嵌套条件计算元组的元素相加：这在元组本身的长度上是线性的，我们将其表示为k。然后整个算法在O(k * |A| * |B|) 中运行。

这是对您当前的O(k * |A| * |B| * |C|) 算法的重大改进。

(2) 矩阵绘图

使用dok_matrix，一种稀疏的 SciPy 矩阵表示。然后你可以在矩阵上使用任何你喜欢的热图绘图库，例如Seaborn's heatmap.