在 numpy 数组之间应用自定义函数/运算符

Question

有两个数组，我想根据已知的单个元素距离获取两个数组之间的距离。

dist = {(4,3): 0.25, (4,1):0.75, (0,0):0, (3,3):0, (2,1):0.25, (1,0): 0.25}
a = np.array([[4, 4, 0], [3, 2, 1]])
b = np.array([[3, 1, 0]])

a
array([[4, 4, 0],
       [3, 2, 1]])
b
array([[3, 1, 0]])

expected output based on dictionary dist:
array([[0.25, 0.75, 0.  ],
       [0.  , 0.25, 0.25]])

所以，如果我们需要不同的元素，我们可以使用a!=b。同样，我想应用以下函数而不是!= -

def get_distance(a, b):
    return dist[(a, b)]

得到上面的预期输出。

我试过np.vectorize(get_distance)(a, b) 并且它有效。但我不确定这是否是以矢量化方式执行上述操作的最佳方法。那么，对于两个 numpy 数组，应用自定义函数/运算符的最佳方式是什么？

Answer 1

不要将您的距离映射存储为dict，而是使用np.array 进行查找（如果大小成为问题，则可能使用稀疏矩阵）。

d = np.zeros((5, 4))
for (x, y), z in dist.items():
    d[x, y] = z

然后，只需索引。

>>> d[a, b]
array([[0.25, 0.75, 0.  ],
       [0.  , 0.25, 0.25]])

---

对于稀疏解决方案（代码几乎相同）：

In [14]: from scipy import sparse

In [15]: d = sparse.dok_matrix((5, 4))

In [16]: for (x, y), z in dist.items():
    ...:     d[x, y] = z
    ...:

In [17]: d[a, b].A
Out[17]:
array([[0.25, 0.75, 0.  ],
       [0.  , 0.25, 0.25]])