使用 numpy 结构比较两个数组可能更快的代码

【问题标题】：Possibly faster code using a numpy construct for comparison of two arrays使用 numpy 结构比较两个数组可能更快的代码
【发布时间】：2020-08-22 16:33:08
【问题描述】：

我有以下代码计算1 <= i < j <= n 的对数，这样xs[i] == ys[j]：

def f(xs, ys):
    s = 0
    for j in range(xs.size):
        s += np.sum(xs[:j] == ys[j])
    return s

这在其他过程中被调用了几次，所以我需要它尽可能快（以一些内存成本）。

数组的大小是> 1e6。

有没有更快的等价物利用一些 numpy 魔法摆脱for 循环？

【问题讨论】：

你可以使用 Numba 吗？它是为这种事情而构建的。
尝试使用np.broadcast_to(xs, (xs.shape[0], xs.shape[0]))得到一个重复xs的方阵，然后np.tril将对角线上的元素清零，然后np.sum(zeroed_matrix == ys)...不确定最后一部分，但是玩弄它
np.sum(zeroed_matrix==ys[:, np.newaxis]) 如this answer

标签： python arrays numpy optimization

【解决方案1】：

如果xs 和ys 大小相同，一种方法：

s = np.triu(xs[:,None]==ys,1).sum()

如果xs 和ys 的大小不同（根据您的代码，您只需将相同长度的ys 与xs 进行比较。如果您想将xs 与所有ys 进行比较，使用上面的行）：

s = np.triu((xs[:,None]==ys[:xs.size]),1).sum()

或等效：

s = (xs[:,None]==ys)[np.triu_indices(xs.size,1)].sum()

将二维xs 的所有元素与ys 进行比较，并将上三角形上的等号相加（与循环的内线相同）

如果您的数组太大并且遇到内存问题，只需将数组分块并使用上面的线条比较对角线上的块加上上三角形上的所有非对角块并将它们相加。

【讨论】：