使用numpy进行矢量化的更好方法

Question

我想知道是否有更好的方法来对以下内容进行矢量化：

array = []
for ele in long_list:
    if condition:
        array.append(1)
vector = np.array(array)

现在除了np.ones(len(long_list))条件外，这基本上是if。

是否有更好的（矢量化）方法来实现上述目标？

Answer 1

正如您在问题中所建议的那样，如果您一次评估整个数组的条件（使用向量化的NumPy操作），而不是在循环中单独检查每个元素，您将获得更好的性能。

使用向量化操作检查条件会创建一个布尔数组，您可以对其进行压缩和强制转换，以生成所需的输出。

为了说明这一点，这里是一个检查数组元素是否小于10的示例：

>>> long_list = np.arange(500)
>>> b = long_list < 10
>>> np.compress(b, b).astype(int)
array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

Answer 2

列表推导比for循环更快。我假设条件取决于ele。

array = [1 for ele in long_list if condition(ele)]

vector = np.array(array)

另一种选择是使用更快的方式来计算condition(ele)为True的情况。但我所知道的那些技巧通常取决于那种情况。

Answer 3

vector = np.fromiter((1 for ele in long_list if condition), dtype=int)

这应该比以前更快：

vector = np.ones(len([x for x in long_list if conditon]))

或

vector = np.array( [1 for ele in long_list if condition(ele)] )

因为(1 for ele in long_list if condition)返回的是生成器而不是列表，所以它避免构建过滤后的列表。

修改

我想知道这是否也是一个可行的选择：

vector = np.array( (1 for ele in long_list if condition(ele)) )