python - 基于复杂索引方案分配 numpy 数组值

Question

我正在尝试使用从第二个 [m x n x f] 数组中提取的值填充第一个 [m x n] numpy 数组，其中索引 f 由第三个 [m x n] 数组提供，但我尚未找到一种无需昂贵的 for 循环即可实现此目的的方法。这是一个简单的代表性示例:

import numpy as np
import random

m, n, k  = 3, 2, 4
np.random.seed(1234)
a = np.random.randint(0,9,(m, n, k))
print(a)
b = np.random.randint(0,k,(m,n))
print(b)
c = np.zeros((m, n))
print(c)

我现在想用 a 中的值填充 c，如下所示:

c[i,j] = a[i,j,b[i,j]]

换句话说，对于 c 中的每个位置 [i,j ]，从 a 中沿前两个轴的相同位置 [i,j] 以及通过评估给出的第三轴的位置处绘制值数组 b 位于 [i,j] 位置。在我的示例中(您的随机整数可能会有所不同？)，数组的值如下:

a :
[[[3 6 5 4]
  [8 1 7 6]]
 [[8 0 5 0]
  [6 2 0 5]]
 [[2 6 3 7]
  [0 0 3 2]]]

b :
[[3 0]
 [1 3]
 [3 3]]

期望的结果是:

c :
[[4 8]
 [0 5]
 [7 2]]

我们得到这个结果如下:从 i, j = 0,0 开始。然后，c[0,0] = [a[0,0,b[0,0]]，并且由于 b[0,0] 的值为 3，因此 a 的计算值为 a[0,0,3]，即 4。继续 i, j = 0, 1。这里 c[0,1] = [a[0,1,b[0,1]]，并且由于 b[0,1] 的值为 0， a 在 a[0,1,0] 处求值，即 8。依此类推。我很清楚如何使用嵌套循环来做到这一点，但我试图避免这种情况，因为我的实际数组要大得多。非常感谢您的帮助!

Answer 1

Numpy 索引广播到输出数组的形状:

c = a[np.arange(b.shape[0])[:, None], np.arange(b.shape[1]), b]

您还可以使用np.take_along_axis以避免在其他轴中手动生成索引，但这里索引必须广播到输入数组:

c = np.take_along_axis(a, b[..., None], -1)