python - 我可以使用数组作为索引对 numpy 数组进行切片吗？

Question

我有 2 个 numpy 数组:

a = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])

b = np.array([2, 1, 2])

我想用b作为 a 列的起始索引并设置 a 的所有值从这些列索引开始到 0 像这样:

np.array([[1, 2, 3],
          [4, 0, 6],
          [0, 0, 0]])

即，将第 1 列的元素从位置 2 开始设置为 0，将第 2 列的元素从位置 1 开始设置为 0，并将第 3 列的元素从位置 2 开始设置为 0。
当我尝试这个时:

a[:, b:] = 0

我得到

TypeError: only integer scalar arrays can be converted to a scalar index

有没有办法在没有 for 循环的情况下使用索引数组进行切片？
编辑:更新示例以显示索引可以是任意的

Answer 1

您可以使用 boolean array indexing .首先，创建要设置为 0 的索引掩码，然后将掩码应用于数组并分配替换值(例如，在您的情况下为 0)。

mask = b>np.arange(a.shape[1])[:,None]
a[~mask]=0

输出:

array([[1, 2, 3],
       [4, 0, 6],
       [0, 0, 0]])