如何向量化这两个嵌套的 for 循环？答案

【问题标题】：How to vectorise these two nested for-loops?如何向量化这两个嵌套的 for 循环？
【发布时间】：2021-09-28 14:25:10
【问题描述】：

我有以下代码：

import numpy as np

my_array = np.zeros((42, 123, 2021))

assert another_array.shape == (42, 123)

for i in range(42):
    for j in range(123):
        my_array[i, j, another_array[i, j]] = 1

假设 another_array 的值保持在正确的范围内（即 another_array 的值是 0 到 2020 之间的整数）。

我想摆脱两个 for 循环。有没有办法对这样的东西进行矢量化？

【问题讨论】：

标签： python numpy vectorization

【解决方案1】：

试试：

my_array = np.zeros((42, 123, 2021))
# assert another_array.shape == (42, 123)
my_array[ np.arange(42)[:,None], np.arange(123), another_array] = 1

这个想法是用一对广播到 (42,123) 的范围替换 i,j 以匹配第 3 轴索引数组。

【讨论】：

这比我使用 .flatten() 的解决方案快很多。谢谢！

【解决方案2】：

我明白了：

import numpy as np

my_array = np.zeros((42, 123, 2021))
my_array = my_array.reshape(-1, 2021)

assert another_array.shape == (42, 123)
another_array = another_array.flatten()

my_array[range(my_array.shape[0]), another_array] = 1
my_array = my_array.reshape(42, 123, 2021)

【讨论】：

for 循环在我的机器上平均快了约 1.14 倍。
比我能找到row, col = np.indices(another_array.shape); my_array[row.ravel(), col.ravel(), another_array.ravel()] = 1 的 for 循环稍快 (~1.12x) 的解决方案。不值得回答。
@MichaelSzczesny 您是否使用足够大的阵列进行测试？众所周知，如果你只使用几个值，numpy 往往比纯 python 慢。
@user2640045 - 我用 OPs 数组和 25 倍的大小进行了基准测试。我认为这对于 OP 的用例会特别有趣。
另一种方法是使其成为一个函数并@numba.jit 它。