假设我有一个数组:
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]])
另一种是使用b
与np.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]])