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

Question

我有以下代码：

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 循环。有没有办法矢量化这样的东西？

Answer 1

我知道了：

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)

Answer 2

尝试：

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 轴索引数组。

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

How to vectorise these two nested for-loops?

python

numpy

vectorization