使用索引数组索引多维 Numpy 数组

Question

假设我有一个 Numpy 数组 a，形状为 (10, 10, 4, 5, 3, 3)，以及两个索引列表，b 和 c，形状分别为 (1000, 6) 和 (1000, 5)，分别代表索引和部分数组的索引。我想使用索引来访问数组，分别生成形状为(1000,) 和(1000, 3) 的数组。

我知道有几种方法可以做到这一点，但它们都很笨拙而且相当不符合 Python 风格，例如将索引转换为元组或分别索引每个轴。

a = np.random.random((10, 10, 4, 5, 3, 3))
b = np.random.randint(3, size=(1000, 6))
c = np.random.randint(3, size=(1000, 5))

# method one

tuple_index_b = [tuple(row) for row in b]
tuple_index_c = [tuple(row) for row in c]

output_b = np.array([a[row] for row in tuple_index_b])
output_c = np.array([a[row] for row in tuple_index_c])

# method two

output_b = a[b[:, 0], b[:, 1], b[:, 2], b[:, 3], b[:, 4], b[:, 5]]
output_c = a[c[:, 0], c[:, 1], c[:, 2], c[:, 3], c[:, 4]]

显然，这些方法都不是很优雅，也不是很容易扩展到更高的维度。第一个也很慢，有两个列表推导，第二个需要你分别写出每个轴。直观的语法 a[b] 出于某种原因返回一个形状为 (1000, 6, 10, 4, 5, 3, 3) 的数组，可能与广播有关。

那么，有没有办法在 Numpy 中做到这一点而不涉及太多的体力劳动/时间？

编辑：并不是真正的重复，因为这个问题涉及多维索引列表，而不仅仅是单个索引，并且产生了一些有用的新方法。

Answer 1

您可以将索引转换为每列是一个单独元素的元组，然后使用__getitem__，假设索引始终是前几个维度：

a.__getitem__(tuple(b.T))

或者简单地说：

a[tuple(b.T)]

---

(a.__getitem__(tuple(b.T)) == output_b).all()
# True

(a.__getitem__(tuple(c.T)) == output_c).all()
# True

(a[tuple(b.T)] == output_b).all()
# True

(a[tuple(c.T)] == output_c).all()
# True

Answer 2

方法三：

output_b = a[map(np.ravel, b.T)]
output_c = a[map(np.ravel, c.T)]