NumPy 通过索引获取多行的快速方法

Question

Given a multidimensional array of shape (6, 100, 2) , and a list of indexes (for example: [1,2,3,1,5,0,0,0,1] ) and I would like to get a numpy ndarray with the data as in these indexes (meaning, the output shape is (9, 100, 2) and rows 6,7,8 are duplicates of each other).

The naive solution:

arr = []
indexes = [1,2,3,1,5,0,0,0,1]
for i in indexes:
  arr.append(data[i])
arr = np.stack(arr)

This naive solution wastes a lot of time in __getitem__ and the stack operation, and I was wondering if there is a faster way to do this with numpy?

Answer 1

You can just use subscripting. So if data is your 6×100×2 matrix, and indexes is your list of indexes, you can retrieve the result with:

result = data