python - 如何在 Python 中基于 2D 数组对 NumPy 3D 数组进行索引？

Question

假设我有一个形状为 (66,5) 的 NumPy 数组 A 和形状为 (100, 66, 5) 的 B。

A 的元素将索引 B 的第一个维度 (axis=0)，其中的值从 0 到 99 (即 B 的第一个维度是 100)。

A = 
array([[   1,    0,    0,    1,    0],
       [   0,    2,    0,    2,    4],
       [   1,    7,    0,    5,    5],
       [   2,    1,    0,    1,    7],
       [   0,    7,    0,    1,    4],
       [   0,    0,    3,    6,    0]
       ....                         ]])

例如，A[4,1] 将取 B 第一维的索引 7，B 第二维的索引 4 和 B 的索引 1第三维B。

我想要的是生成形状为 (66,5) 的数组 C，其中包含 B 中的元素，这些元素是根据 中的元素选择的>A.

Answer 1

您可以使用 np.take_along_axis这样做:

import numpy as np
np.random.seed(0)
a = np.random.randint(100, size=(66, 5))
b = np.random.random(size=(100, 66, 5))
c = np.take_along_axis(b, a[np.newaxis], axis=0)[0]
# Test some element
print(c[25, 3] == b[a[25, 3], 25, 3])
# True