python - 如何向量化这两个嵌套的 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

尝试:

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