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

Question

我正在尝试使用从第二个 [mxnxf] 数组中提取的值填充第一个 [mxn] numpy 数组，其中索引 f 由第三个 [mxn] 数组提供，我还没有找到一种无需昂贵的 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)

Answer 2

使用np.unravel_index创建索引，索引然后.reshape：

c = np.empty((m, n), dtype=np.int32)
ro, co = np.unravel_index(np.arange(n * m), (m, n))
c[:, :] = a[ro, co, b[ro, co]].reshape(3,2)
print(c)

输出

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