Question

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

我知道有几种方法可以做到这一点，但它们都是笨拙的而且不是pythonic的，例如将索引转换为元组或分别索引每个轴。

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)]

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

2 个答案: