Python: 多维数组掩码

Question

下面这段代码在Matlab中是这样的，我们想知道在Python中怎么写才能达到同样的效果。

Matlab:

B = 2D array of integers as indices [1...100]
A = 2D array of numbers: [10x10]
A[B] = 0

这段代码运行得很好，比如当你写 B[i]=42 时，它能找到第 5 列的第 2 行来进行设置。

但是在Python中，如果这样写就会出现错误：超出范围，这也是合理的。不过我们想把上面的Matlab代码转换成Python，所以需要找到更符合Python风格的方法。

另外，我们还要考虑更高维度的问题，比如：

B = 2D array of integers as indices [1...3000]
C = 3D array of numbers: [10x10x30]
C[B] = 0

我们想到的一种方法是把索引数组的元素改成 i,j 这样的形式，而不是绝对位置。也就是说，把位置 42 转换成 divmod(42,m=10)[::-1] >>> (2,4)。这样我们就可以得到一个 nx2 >>> ii,jj 的索引向量，这样在索引 A 时会更方便。

我们觉得这种方法可能更好，而且在Python中对于更高维度的情况也会更高效。

Answer 1

你可以在对数组（A）进行索引之前，先用 .ravel() 把它变成一维的，然后再用 .reshape() 来调整它的形状。

另外，因为你知道 A.shape 的形状信息，你也可以在对另一个数组（B）进行索引之前，使用 np.unravel_index。

示例 1:

>>> import numpy as np
>>> A = np.ones((5,5), dtype=int)
>>> B = [1, 3, 7, 23]
>>> A
array([[1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1]])
>>> A_ = A.ravel()
>>> A_[B] = 0
>>> A_.reshape(A.shape)
array([[1, 0, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1]])

示例 2:

>>> b_row, b_col = np.vstack([np.unravel_index(b, A.shape) for b in B]).T
>>> A[b_row, b_col] = 0
>>> A
array([[1, 0, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1]])

---

后来发现：你还可以使用 numpy.put

>>> import numpy as np
>>> A = np.ones((5,5), dtype=int)
>>> B = [1, 3, 7, 23]
>>> A.put(B, [0]*len(B))
>>> A
array([[1, 0, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 0, 1]])