替换索引数组下方的numpy 2D数组元素

Question

我搜索了之前的问题，并没有找到一个非常适合这个问题。

我将A的第j列中的所有单元格的值设置为0，其行的行数小于m [j]：

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'循环的情况下重写该脚本？

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相同形状的布尔掩码，然后A其用于将值设置为0。

注意 - 如果A保证是一个数组，我会推荐Divakar的解决方案 ，这与此非常相似。

Answer 2

对长度为n的远程数组使用外部广播比较直接获得A -

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

为了提高性能，我们可以使用类型转换来降低精度，例如.astype(np.int)甚至.astype(np.uint8) 。 如果您对布尔数组没问题，请完全跳过类型转换。

对于int dtypes，我们也可以简单地view它，这为我们节省了一些内存 -

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