我有一个带有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')
我想使用列索引ind
将bar
的值分配给矩阵foo
。我想每行都这样做。假设执行此操作的函数称为assign_function
,我的输出将如下所示:
>>> assign_function(ind, bar, foo)
>>> print(foo)
array([[9, 1],
[8, 3],
[4, 7]])
有没有pythonic的方法来做到这一点?
答案 0 :(得分:4)
由于ind
处理第一个轴,因此只需要第0个轴的索引器。您只需使用np.arange
:
foo[np.arange(len(foo)), ind] = bar
foo
array([[9, 1],
[8, 3],
[4, 7]])
答案 1 :(得分:3)
将broadcasting
与masking
-
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]])