使用两个 2D 数组索引 4D NumPy 数组

Question

我有一个 4D 目标 NumPy 数组，我想使用 2D 源数组中的值填充它，使用两个额外的 2D 数组，这两个数组指定目标数组第二和第三轴中的位置，其中来自源的值应放置阵列。下面带有一些示例值的代码可以使用 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]

我如何向量化这种方法？我知道如果目标数组是二维的，源数组和索引数组都是一维的，我可以使用 this approach，但我不确定如何将它推广到我的情况。

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]

由于您提供了代码，但没有示例，我将自己跳过该步骤，'eyeball' 一个答案 - 无需测试。

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

想法是为 t 和 d 创建索引数组，使用 (T,D) 形数组进行广播。

使用两个 2D 数组索引 4D NumPy 数组

Indexing a 4D NumPy Array with two 2D arrays

python

arrays

numpy

multidimensional-array

matrix-indexing