主页 / coding / 问题 / 79337016
Accepted
Kuraga
Asked: 2025-01-08 02:37:00 +0800 CST

从另一个数组投影一个索引

  • 772

a我有一个大小为 的数组(M, N, K)。我有一个b大小为 的数组(M, N)，其整数值为[0，K-1]。

如何以最简单的方式 获取c大小为 的数组 ？(M, N)c[i, j] == a[i, j, b[i, j]]

它是索引指南的哪一部分？

numpy

2 个回答

  1. 要通过这种方式提取数据，您需要将相应的索引作为 3 个维度中的 2 维数组传递，其中前 2 个应该是meshgrid矩阵的结果（注意indexign='ij'）：

    a[*np.meshgrid(range(M), range(N), indexing='ij'), b]

    有关详细信息，请参阅NumPy 和numpy.meshgrid中的高级索引。

    例子：

    import numpy as np

M, N, K = 3, 4, 5
a = np.arange(M*N*K).reshape(M, N, K)
b = np.arange(M*N).reshape(M, N) % K
first, second = np.meshgrid(range(M), range(N), indexing='ij')
c = a[first, second, b]

assert all(a[i, j, b[i,j]] == c[i,j] for i in range(M) for j in range(N))
    • 2
  2. Best Answer

    您可以使用高级索引：

    c = a[np.arange(M)[:, None], np.arange(N), b]

    输出：

    array([[ 0,  6, 12, 18],
       [24, 25, 31, 37],
       [43, 49, 50, 56]])

    以@Vitalizzare 为例：

    # input
M, N, K = 3, 4, 5
a = np.arange(M*N*K).reshape(M, N, K)
b = np.arange(M*N).reshape(M, N) % K

# output
array([[ 0,  6, 12, 18],
       [24, 25, 31, 37],
       [43, 49, 50, 56]])
    • 2

相关问题