如何用numpy替换double for循环？答案

【问题标题】：How to replace double for loop with numpy?如何用numpy替换double for循环？
【发布时间】：2023-02-26 01:40:19
【问题描述】：

假设我有一个列表。然后我必须遍历该列表中的每个元素并将其与该列表中的每个其他元素进行比较。基于这种比较，我需要将这个元素添加到另一个列表中。目前，我用 for 循环解决这个问题。但我想通过使用 numpy 的矢量化来优化它。

我目前的解决方案如下：

a = [1, 2, 3, 4, 5]
b = []
for x in a:
  for y in a:
    if comparison_function(x, y):
      b.append((x, y))

comparison_function 只是根据某些标准比较这两个元素。显然，这不是一个很好的解决方案，我想通过矢量化来解决这个问题。如何实现这一目标？

【问题讨论】：

需要将此元素添加到另一个列表。- 你附加了一个元组(x,y)，它是作为比较的结果还是元素x,y原样？ comparison_function 本质上是什么？
比较函数中的顺序重要吗？如果不是，您可以将运行时间减半。
没有通用的方法。您需要对其进行矢量化的是 comparison_function。
这个问题没有具体的comparison_function是没有意义的
你可以简单地使用理解：b = [(x,y) for x in a for y in a if comparison_function(x, y)]

标签： python numpy vectorization

【解决方案1】：

这是一个使用“小于”作为比较函数的玩具示例：

import numpy as np

a = np.array([1, 2, 3, 4, 5])

您可以使用广播来计算所有对的比较：

m = a[:, None] < a[None, :]

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

最后，您可以使用 np.argwhere 生成 True 值的索引，您可以使用这些索引获取 a 中的原始值：

out = a[ np.argwhere(m) ]

出去：

array([[1, 2],
       [1, 3],
       [1, 4],
       [1, 5],
       [2, 3],
       [2, 4],
       [2, 5],
       [3, 4],
       [3, 5],
       [4, 5]])

您可以选择将其转换回 list[tuple[int, int]]，但您也可以将其保留为 (N, 2) 形状的数组。

【讨论】：