用索引数组替换numpy 2D数组中某些元素

Question

我搜索了之前的问题，没有一个完全符合这个问题。

我正在将A矩阵中第j列所有行下标小于m[j]的单元格值设为0：

import numpy as np
n = 8
A = np.ones([n,n])
m = np.array([1,1,1,2,2,2,3,4])
for j in range(n):
    for i in range(m[j]):
        A[i,j] = 0

如何重写没有“for”循环的脚本？

- Forklift17

Answer 1

您可以使用广播的NumPy比较-

>>> A[np.arange(n)[:, None] < m] = 0
>>> 
array([[0., 0., 0., 0., 0., 0., 0., 0.],
       [1., 1., 1., 0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1., 1., 0., 0.],
       [1., 1., 1., 1., 1., 1., 1., 0.],
       [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., 1., 1., 1., 1., 1., 1., 1.]])

在此，调用[:, None]会增加np.arange(n)的形状，使得比较运算符<广播到范围内每个元素的m。这将生成一个与A相同形状的布尔掩码，然后用于将值设置为0。

注意 - 如果可以保证A是一个全为1的数组，我建议使用Divakar's solution，这与本方法非常相似。

- cs95

Answer 2

使用针对长度为n的范围数组的外部广播比较，以直接获取A -

A = (m <=  np.arange(n)[:,None]).astype(float)

为了提高性能，我们可以使用类型转换为较低精度的类型，例如.astype(np.int)甚至.astype(np.uint8)。如果您可以接受布尔数组，则可以完全跳过类型转换。

对于整数数据类型，我们也可以简单地view到它，这样可以节省一些内存。

(m <=  np.arange(n)[:,None]).view(np.int8)

- Divakar

---

如果一开始输入的全是1，那么像这样初始化A是有意义的。+1 - cs95