从基于 2D 矩阵位置的 3D 矩阵获取值

Question

假设我有一个 numpy 矩阵：data = np.random.rand(200, 50, 100) 并且有我需要的位置：locs = np.random.randint(50, size=(200, 2)).

我如何获得形状为 (200, 2, 100) 的结果矩阵？本质上，我想在 locs 指定的位置从 data 获取值。

如果我这样做：data[locs]，我最终会得到形状为 (200, 2, 50, 100) 而不是 (200, 2, 100).

的结果矩阵

根据要求更新了更多详细信息：

如果我们有：

data = np.arange(125)
reshaped = np.reshape(data, (5, 5, 5))
locs = [[3, 4], [2, 1], [1, 3], [3, 3], [0, 0]]

然后执行类似 data[locs] 的操作应该会产生以下输出：

array([[[ 15,  16,  17,  18,  19],
        [ 20,  21,  22,  23,  24]],

       [[ 35,  36,  37,  38,  39],
        [ 30,  31,  32,  33,  34]],

       [[ 55,  56,  57,  58,  59],
        [ 65,  66,  67,  68,  69]],

       [[ 90,  91,  92,  93,  94],
        [ 90,  91,  92,  93,  94]],

       [[100, 101, 102, 103, 104],
        [100, 101, 102, 103, 104]]])

Answer 1

高级索引的结果将是沿着您正在索引的维度的索引形状。 data[locs] 等同于 data[locs, :, :]，所以你的形状将是 locs.shape + data.shape[1:]，或 (200, 2, 50, 100)。

你出现的要求是使用locs索引data的轴1，保持轴0与[中的行同步=16=]。为此，您需要使用 locs 沿轴 1 进行索引，并在轴 0 中提供从 0 到 200 的索引。

重要的是要记住所有高级索引必须广播到相同的形状。由于 locs 的形状为 (200, 2)，第一个索引的形状必须为 (200, 1) 或 (200, 2) 才能正确广播。我将展示前者，因为它更简单、更高效。

data = np.random.rand(200, 50, 100)
locs = np.random.randint(50, size=(200, 2))
rows = np.arange(200).reshape(-1, 1)

result = data[rows, locs, :]

从基于 2D 矩阵位置的 3D 矩阵获取值

Obtaining values from a 3D matrix based locations in a 2D matrix

python

indexing

numpy

array-broadcasting