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

时间:2019-02-02 16:21:19

标签: python numpy vectorization

假设我有一个数组:

rbind_list(lapply(list1, function(x) setNames(data.frame(t(x$values)), x$names)))

#     a     b     c     x
#   <dbl> <dbl> <dbl> <dbl>
#1    25    13    11    NA
#2    12    10    NA     2

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

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

实现这一目标的最麻木的方法是什么?我知道我可以使用简单的for循环来做到这一点,但这对numpy来说是不好的做法。我的直觉是,可能有一些聪明的快速方法可以通过广播和矢量化来做到这一点,但我无法弄清楚。

1 个答案:

答案 0 :(得分:1)

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

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]])