时间戳

首页 > 解决方案 > 优化测试来自多个 NumPy 数组的所有行组合

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

问题描述

我有三个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）我可以使用嵌套循环编写上述内容，但有没有更快的方法，也许是列表推导或itertools？

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

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

编辑：评论中要求我的丑陋代码

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)

标签： pythonarraysnumpymatrixoptimization

解决方案

我们可以利用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

推荐阅读