我搜索了之前的问题,没有一个完全符合这个问题。
我正在将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”循环的脚本?
>>> 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,这与本方法非常相似。使用针对长度为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)