python - 为具有指定列的每一行分配值到一个 numpy 数组

Question

我有一个矩阵 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]])

我有一个包含 n 元素的数组 bar。示例:

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

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

>>> 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 的方法来做到这一点？

Answer 1

由于 ind 处理第一个轴，您只需要第零个轴的索引器。您可以使用 np.arange 非常简单地完成此操作:

foo[np.arange(len(foo)), ind] = bar
foo

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