是否可以通过另一个二维数组逐行索引二维数组而无需在 NumPy 中循环？

假设我有带索引的 NumPy 数组：

array([[1, 3],
       [2, 5]])

另外，我有带零的 NumPy 数组：

array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]])

我需要从带有索引的数组中取出第一行，并在数组的第一行中为这些索引设置 1 为零。此后第一行必须看起来像[0, 1, 0, 1, 0, 0].

然后我还需要为数组中的第二行做零，并分别使用数组中的第二行和索引。在这之后的第二行必须看起来像[0, 0, 1, 0, 0, 1].

所以，我想得到这样的数组：

array([[0, 1, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1]])

我怎么能在没有任何循环的情况下做到这一点？

回答

>>> i = np.array([[1, 3],
[2, 5]])
>>> x = np.array([[0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0.]])

>>> np.put_along_axis(x, i, values=1, axis=1)
>>> x
array([[0., 1., 0., 1., 0., 0.],
[0., 0., 1., 0., 0., 1.]])