如何使用不同形状的索引数组切片Numpy数组？

假设我们有如下二维Numpy数组：

import numpy as np

arr = np.array([[1,1,0,1,1],
                [0,0,0,1,0],
                [1,0,0,0,0],
                [0,0,1,0,0],
                [0,1,0,0,0]])

以及指定的行和列索引数组：

rows = np.array([0,2,4])
cols = np.array([1,2])

我们的目标是用这两个索引数组对arr切片，得到如下结果：

arr_sliced = np.array([[1,0],
                       [0,0],
                       [1,0]])

但如果直接执行arr[rows, cols]进行索引，会触发报错：

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (3,) (2,)

解决方案 #

解决这个问题的核心是利用Numpy的广播机制，只需对行索引数组增加一个维度，让两个索引数组的形状满足广播规则即可：

arr[rows[:, np.newaxis], cols]
# 转换后 rows[:, np.newaxis].shape == (3,1)
# cols.shape == (2,)

原理说明 #

将rows从形状(3,)转换为(3,1)后，它就能和形状为(2,)的cols触发广播机制，最终生成(3,2)的索引组合，正好匹配我们需要的切片结果。

内容的提问来源于stack exchange，提问作者lezaf