使用两个 2D arrays 索引 4D NumPy 阵列

Question

我有一个 4D 目标 NumPy 数组，我想用 2D 源数组中的值填充它，使用两个额外的 2D arrays 指定 position 的第二个数组和第三个源数组的值。 下面带有一些示例值的代码可以使用 for 循环方法执行此操作：

import numpy as np

T, M, N, D = 5, 4, 3, 2  # Example dimensions
target = np.zeros((T, M, N, D))  # Target array
source = np.arange(T*D).reshape(T, D)  # Example source array with arbitrary numbers
index_dim_1 = np.array([  # Array of shape T, D with values between 0 and M-1
    [0, 3],
    [1, 2],
    [2, 1],
    [3, 0],
    [0, 3]
])
index_dim_2 = np.array([  # Array of shape T, D with values between 0 and N-1
    [0, 2],
    [1, 1],
    [2, 0],
    [0, 2],
    [1, 1]
])

for t in range(T):
    for d in range(D):
        m = index_dim_1[t, d]
        n = index_dim_2[t, d]
        target[t, m, n, d] = source[t, d]

我怎样才能矢量化这种方法？ 我知道目标数组是否是二维的，源和索引 arrays 都是一维的，我可以使用这种方法，但我不确定如何将其推广到我的情况。

Answer 1

for t in range(T):
    for d in range(D):
        n = index_dim_1[t, d]
        m = index_dim_2[t, d]
        target[t, n, m, d] = source[t, d]

由于您提供了代码，但没有提供示例，因此我将自己跳过该步骤，并“目测”答案-无需测试。

 target[np.arange(T)[:,None], index_dim_1, index_dim_2, np.arange(D)] = source

这个想法是为t和d创建索引 arrays 以 (T,D) 形状的 arrays 广播。