从 3D Numpy 数组中使用多维索引提取元素：使用案例和通用的解决方案

2024-03-17 01:00:33

从 3D Numpy 数组中使用多维索引提取元素

介绍

在处理多维数组时，提取特定元素是一项常见的任务。在本文中，我们将探究如何从 3D Numpy 数组中使用两个一维 Numpy 数组作为列和深度索引来提取元素。我们将通过一个示例来演示这个过程，并在其后提供一个更通用的解决方案。

问题

假设我们有一个 3D Numpy 数组 A（形状：NxMxK），以及两个一维整型数组 B 和 C（形状：1xN）。数组 B 的值介于 0 到 M-1 之间（包括端点），而数组 C 的值介于 0 到 K-1 之间（包括端点）。我们的目标是返回 3D 数组 A[n,m,k] 的值，其中对于 B 和 C 具有相同形状的每个元素索引 h，m 是 B 中索引 h 处的值，而 k 是 C 中索引 h 处的值。

示例

为了说明该过程，我们使用以下示例数组和索引数组：

import numpy as np

N, M, K = 3, 2, 4
A = np.random.uniform(100, 120, 3 * 2 * 4).reshape(N, M, K)  # 随机生成的 3D 数组
B = np.random.randint(0, M, N)  # 随机生成的索引数组 1
C = np.random.randint(0, K, N)  # 随机生成的索引数组 2

为了使用数组 B 作为列索引和数组 C 作为深度索引从数组 A 中提取元素，我们可以使用以下步骤：

将数组 B 和 C 转换为与数组 A 相同的形状。我们可以使用 numpy.broadcast_to 函数来执行此操作：

B_broadcast = np.broadcast_to(B[:, None], A.shape[:2])
C_broadcast = np.broadcast_to(C[:, None, None], A.shape)

使用 np.take_along_axis 函数沿轴 1（列轴）提取数组 A 中使用数组 B 作为索引的值：

D = np.take_along_axis(A, B_broadcast, axis=1)

使用 np.take_along_axis 函数沿轴 2（深度轴）提取数组 A 中使用数组 C 作为索引的值：

D = np.take_along_axis(D, C_broadcast, axis=2)

数组 D 将包含从数组 A 中提取的元素，其中 m 是数组 B 中的列索引，而 k 是数组 C 中的深度索引。

结论

在本文中，我们探索了如何从 3D Numpy 数组中使用两个一维 Numpy 整型数组作为列和深度索引来提取元素。我们通过一个示例演示了该过程，并提供了通用的解决方案，该解决方案利用了 np.take_along_axis 函数。通过理解这些技术，您可以有效地在多维数组中定位和提取数据，从而增强您的 Numpy 技能。