如何根据随机索引数组将向量中的元素插入到矩阵中

Question

基本上，我试图根据随机索引将向量中的元素插入到矩阵中

size = 100000
answer_count = 4
num_range = int(1e4)

a = torch.randint(-num_range, num_range, size=(size, ))
b = torch.randint(-num_range, num_range, size=(size, ))
answers = torch.randint(-num_range, num_range, size=(size, answer_count))

for i in range(size): answers[i, np.random.randint(answer_count)] = a[i] + b[i]

我试过

c = a + b
pos = torch.randint(answer_count, size=(size, ))
answers[:, pos] = c

但我确实做错了什么

Answer 1

我认为你需要像这样更改最后一行：

answers[np.arange(size), pos] = c

问题在于高级索引的使用不正确。要了解这些索引的区别，请尝试打印出 answers[:, pos] 与 answers[np.arange(size), pos]，您会明白为什么前一个不起作用。 answers[np.arange(size), pos] 选择每个 pos 一行 而 answers[:, pos] 选择 所有行 每个 pos. numpy doc here.

中有关高级索引的更多信息