在 numpy 中计算布尔矩阵乘法的最佳方法

Question

我想计算布尔数组与其转置的矩阵乘积：

import numpy as np
a = np.array([[1, 0, 1], [1, 1, 0]], dtype=bool)

最好/最快的方法是什么？

我做了几次尝试：

---

尝试1

out1 = np.matmul(a, a.T)
print(out1)

[[真真][真真]]

仅应用

np.matmul

是行不通的，因为溢出。

---

尝试2

所以我想出了将

np.uint64

输出数组传递给函数的想法：也不起作用。看起来总和是用 bool 计算的，然后转换为 uint64。

out2 = np.zeros((2, 2), dtype=np.uint64)
np.matmul(a, a.T, out=out2)
print(out2)

[[1 1] [1 1]]

---

尝试3

手动计算行*列乘积然后调用 sum 是可行的：但是它占用了太多内存，因为它在求和之前存储了所有乘积结果。

out3 = (a[None,:, :] * a[:, None, :]).sum(axis=-1)
print(out3)

[[2 1] [1 2]]

---

尝试4

所以我尝试在计算之前转换数组：也有效。但是转换输入数组需要 64 倍的内存，我猜想与两个布尔值相比，如果使用 uint64 进行乘法则需要更多时间。

out4 = np.matmul(a.astype(np.uint64), a.T)
print(out4)

[[2 1] [1 2]]

---

我可能已经监督过这个问题，有更好的解决方案吗？

Answer 1

您的计算中没有“溢出”。乘法和加法对于布尔值有不同的含义。

让我们假设另一个例子：

a = np.array([[1, 0, 1],
              [0, 1, 0]], dtype=bool)

out1 = np.matmul(a, a.T)

array([[ True, False],
       [False,  True]])

结果是正确的，如果至少有一个产品不为空，则结果为 True。假设 i,j 为坐标，这意味着：

out[i,j] = a[i,0]*a[j,0] + a[i,1]*a[j,1] + a[i,2]*a[j,2]

带书立 布尔值，

*

相当于

&

，

+

相当于

|

。

# product / & truth table 
np.array([False, True]) * np.array([[False], [True]])

array([[False, False],
       [False,  True]])

# sum / | truth table
np.array([False, True]) + np.array([[False], [True]])

array([[False,  True],
       [ True,  True]])

事实上，如果您首先转换为整数，您将得到相同的结果，除了整数计算：

np.matmul(a.astype(np.uint64), a.T)

array([[2, 0],
       [0, 1]], dtype=uint64)

现在的问题是：你需要什么作为输出？

如果需要整数，则先转换为整数再进行计算。如果你想要一个布尔值，那么原来的方法就可以了。