Numpy:如何为不同索引数组的每一列设置column[index:]=value



假设我有一个数组:

my_array = np.random.normal(size=(5,3))
print(my_array)
[[ 0.45110035 -1.08385534  1.2126054 ]
[ 1.51280316  0.4308235  -0.31839059]
[-0.00348102 -0.50814392  0.00734745]
[-0.63701191  0.95413945 -1.40480595]
[-1.66723431 -0.52822503 -1.14282036]]

以及与列数长度相同的索引列表:

my_indices = np.array([3, 1, 2])

对于my_array中的每一列,我希望从my_indices获得相应的索引,并将该元素和该列的所有后续元素设置为某个新值。换句话说,结果看起来像:

my_new_array = np.something(my_array, my_indices, 23)
print(my_new_array)
[[ 0.45110035 -1.08385534  1.2126054 ]
[ 1.51280316  23.0       -0.31839059]
[-0.00348102  23.0        23.0      ]
[ 23.0        23.0        23.0      ]
[ 23.0        23.0        23.0      ]]
# i.e.:
# col[3:]=23  col[1:]=23  col[2:]=23

实现这一点最愚蠢的方法是什么?我知道我可以用一个简单的for循环来完成,但这对numpy来说是不好的做法。我的直觉是,可能有一些聪明的快速方法可以通过广播和矢量化来做到这一点,但我无法理解。

设置或从每列中选择一项很容易:

In [10]: arr = np.ones((5,3),int)
In [11]: b = np.array([3,1,2])
In [12]: arr[b, np.arange(3)] = 0
In [13]: arr
Out[13]: 
array([[1, 1, 1],
[1, 0, 1],
[1, 1, 0],
[0, 1, 1],
[1, 1, 1]])

但是,为每列设置一个切片更为棘手。最简单的方法就是对列进行迭代。

In [14]: for i,j in enumerate(b):
...:     arr[j:, i] = 0
...:     
In [15]: arr
Out[15]: 
array([[1, 1, 1],
[1, 0, 1],
[1, 0, 0],
[0, 0, 0],
[0, 0, 0]])

另一种是使用bnp.arange(5)的比较来设置掩模。这更快,但没有那么直观。我每次都要重新思考这个方法。

In [16]: mask = np.arange(5)[:,None]>=b
In [17]: mask
Out[17]: 
array([[False, False, False],
[False,  True, False],
[False,  True,  True],
[ True,  True,  True],
[ True,  True,  True]])
In [18]: arr[mask] = 2
In [19]: arr
Out[19]: 
array([[1, 1, 1],
[1, 2, 1],
[1, 2, 2],
[2, 2, 2],
[2, 2, 2]])

最新更新