我可以使用数组作为索引对 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]])

Answer 2

我认为问题出在a[:,b:]；如果b 不是缩放器，这里b: 意义不大，例如5: 表示从第 6 位开始，但 [1,2,3]: 在数组为 2d 时没有任何意义。

应该是a[:,b]。设置a[:,b] = 0 会将b 中指定的所有列设置为0。以下是运行。

In [2]: import numpy as np

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

In [4]: a
Out[4]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [5]: b
Out[5]: array([2, 1, 2])

In [6]: b.dtype
Out[6]: dtype('int64')

In [7]: a[:, b:] = 0
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-6e5050513225> in <module>
----> 1 a[:, b:] = 0

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

In [8]: a[:, b] = 0

In [9]: a
Out[9]: 
array([[1, 0, 0],
       [4, 0, 0],
       [7, 0, 0]])

但这不是你想要的。

要得到你想要的，你需要指定行索引和列索引，例如(1,1), (2,0), (2,1), (2,2).

In [11]: a[[1,2,2,2], [1, 0, 1,2]] = 0

In [12]: a
Out[12]: 
array([[1, 2, 3],
       [4, 0, 6],
       [0, 0, 0]])