时间戳

首页 > 解决方案 > 如何在 Python 中基于 2D 数组对 NumPy 3D 数组进行索引？

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

问题描述

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

的元素A将索引 的第一个维度 ( axis=0) B，其中值是从 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] 将采用 的第一个维度的索引 7、B的第二个维度的索引 4B和第三个维度的索引 1 B。

我想要的是生成C形状数组 (66,5)，其中包含B根据A.

标签： pythonarraysnumpy

解决方案

你可以用它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

推荐阅读