用两个 2D 数组索引 4D NumPy 数组答案

【问题标题】：Indexing a 4D NumPy Array with two 2D arrays用两个 2D 数组索引 4D NumPy 数组
【发布时间】：2022-01-09 07:36:44
【问题描述】：

我有一个 4D 目标 NumPy 数组，我想用来自 2D 源数组的值填充它，使用两个额外的 2D 数组指定目标数组的第二和第三轴上的位置，源数组的值应该在该位置放置。下面带有一些示例值的代码可以使用 for-loop 方法做到这一点：

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，但我不确定如何将其推广到我的情况。

【问题讨论】：

您能否编辑代码以包含一些实际数字。多维数组很容易出错，所以我想确保我匹配你想要做的事情。

标签： python arrays numpy multidimensional-array matrix-indexing

【解决方案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 创建索引数组，使用 (T,D) 形状的数组进行广播。

【讨论】：