为具有指定列的每一行的numpy数组赋值

时间:2017-12-19 16:26:21

标签: python arrays numpy

我有一个带有foo行和n列的矩阵m。例如:

>>> import numpy as np
>>> foo = np.arange(6).reshape(3, 2) # n=3 and m=2 in our example
>>> print(foo)
array([[0, 1],
       [2, 3],
       [4, 5]])

我有一个带有bar元素的数组n。例如:

>>> bar = np.array([9, 8, 7])

我有一个长度为ind的列表n,其中包含列索引。例如:

>>> ind = np.array([0, 0, 1], dtype='i')

我想使用列索引indbar的值分配给矩阵foo。我想每行都这样做。假设执行此操作的函数称为assign_function,我的输出将如下所示:

>>> assign_function(ind, bar, foo)
>>> print(foo)
array([[9, 1],
       [8, 3],
       [4, 7]])

有没有pythonic的方法来做到这一点?

2 个答案:

答案 0 :(得分:4)

由于ind处理第一个轴,因此只需要第0个轴的索引器。您只需使用np.arange

即可轻松完成此操作
foo[np.arange(len(foo)), ind] = bar
foo

array([[9, 1],
       [8, 3],
       [4, 7]])

答案 1 :(得分:3)

broadcastingmasking -

一起使用
foo[ind[:,None] == range(foo.shape[1])] = bar

逐步运行示例 -

# Input array
In [118]: foo
Out[118]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

# Mask of places to be assigned
In [119]: ind[:,None] == range(foo.shape[1])
Out[119]: 
array([[ True, False],
       [ True, False],
       [False,  True]], dtype=bool)

# Assign values off bar
In [120]: foo[ind[:,None] == range(foo.shape[1])] = bar

# Verify
In [121]: foo
Out[121]: 
array([[9, 1],
       [8, 3],
       [4, 7]])