使用两个 1d numpy 整数数组作为列和深度索引从 3d numpy 数组中切片元素

Question

我有一个 3d numpy 数组 A（形状：NxMxK）、两个 1d 整数数组 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) #generates 24 random nubmers in [100,120) and shapes it as a 3d array of 3x2x4.
B=np.random.randint(0, M, N) #generates N=3 random integers in [0,M=2)
C=np.random.randint(0, K, N)#generates N=3 random integers in [0,K=4)

# I want to use B as the index m and C as index k of A so the an array D = np.array([ A[0, B[0], C[0]],A[1, B[1], C[1]], ..., A[N-1, B[N-1], C[N-1]]]) is returned.

Answer 1

试试这个：

import numpy as np

N, M, K = 3, 2, 4

# Generating A
A = np.random.uniform(100, 120, (N, M, K))

# Generating B and C
B = np.random.randint(0, M, N)
C = np.random.randint(0, K, N)

# Creating D using array indexing
D = A[np.arange(N), B, C]

print(D)

输出：

[101.22545738 105.11929713 105.60995168]