将索引列表转换为numpy数组

Question

我有一个大小为 nxd 的 numpy 零数组。我的任务是把这个数组每一行的指定列变成 1。为此，我得到了一个大小为 n 的列表，这个列表的第 i 个值是要变成 1 的索引。

这个任务可以通过下面的 for 循环来完成：

import numpy as np

N=5; D =3
array = np.zeros(shape=(N,D))
ones_index = [0,2,1,0,1]
for row, column in enumerate(ones_index):
    array[row,column] = 1

虽然这样做没问题，但我想 numpy 里应该有更简单的方法来实现这个功能。 有没有 numpy 函数可以把一组索引在数组中转换成特定的值呢？

Answer 1

另一个选择是使用 np.identity：

np.identity(max(ones_index) + 1)[ones_index]

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

Answer 2

一个可能的解决方案：

array = np.zeros(shape=(N, D))

ones_index = [0, 2, 1, 0, 1]
array[np.arange(N), ones_index] = 1

print(array)

输出结果：

[[1. 0. 0.]
 [0. 0. 1.]
 [0. 1. 0.]
 [1. 0. 0.]
 [0. 1. 0.]]

---

编辑：快速对比一下这个答案和Onaymbu的答案：

import numpy as np
import perfplot


def get_identity_andrej(N, D, ones_index):
    array = np.zeros(shape=(N, D))
    array[np.arange(N), ones_index] = 1
    return array


def get_identity_onyambu(N, D, ones_index):
    return np.identity(max(ones_index) + 1)[ones_index]


np.random.seed(0)

perfplot.show(
    setup=lambda n: (n * 5, n * 3, np.random.randint(0, n * 3, size=n * 5)),
    kernels=[
        get_identity_andrej,
        get_identity_onyambu,
    ],
    labels=["andrej", "onyambu"],
    n_range=[1, 2, 3, 5, 10, 100, 250, 500, 1000],
    xlabel="N",
    logx=True,
    logy=True,
    equality_check=None,
)

在我的电脑上（AMD 5700x），我得到了这个结果：