索引多维numpy的数组索引

Question

数组

让我们说我有一个numpy的阵列a与形状(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)]