如何在给定每行一维索引的二维 numpy 数组中设置值？

Question

在 numpy 中，您可以将一维数组的索引设置为一个值

import numpy as np 

b = np.array([0, 0, 0, 0, 0])

indices = [1, 3]

b[indices] = 1
b

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

我正在尝试以最编程优雅和计算效率最高的方式使用多行和每行的索引来做到这一点。 例如

b = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])

indices = [[1, 3], [0, 1], [0, 3]]

期望的结果是

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

我尝试了b[indices]和b[:,indices]但它们导致错误或不希望的结果。

通过搜索，有一些变通方法，但每个都倾向于在 python 中至少需要 1 个循环。

解决方案 1：循环遍历 2d 数组的每一行。 这样做的缺点是循环在 python 中运行，这部分不会利用 numpy 的 c 处理。

解决方案 2：使用 numpy put 。 缺点是输入数组的扁平化版本，因此索引也需要扁平化，并根据行大小和行数进行更改，这put在 python 中使用双循环。

解决方案 3： put_along_axis似乎每行只能设置 1 个值，因此我需要对每行的值数重复此函数。

什么是计算和编程上最优雅的解决方案？ 任何 numpy 可以处理所有操作的地方？

Answer 1

在 [330] 中： b = np.zeros((3,5),int)

要设置 (3,2) 列，行索引需要是 (3,1) 形状（通过广播匹配）：

In [331]: indices = np.array([[1,3],[0,1],[0,3]])

In [332]: b[np.arange(3)[:,None], indices] = 1

In [333]: b
Out[333]: 
array([[0, 1, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 0]])

put along做同样的事情：

In [335]: b = np.zeros((3,5),int)
In [337]: np.put_along_axis(b, indices,1,axis=1)

In [338]: b
Out[338]: 
array([[0, 1, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 0]])

Answer 2

关于在每个维度中构建索引然后使用基本索引的解决方案：

from itertools import chain

b = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])

# Find the indices along the axis 0
y = np.arange(len(indices)).repeat(np.fromiter(map(len, indices), dtype=np.int_))

# Flatten the list and convert it to an array
x = np.fromiter(chain.from_iterable(indices), dtype=np.int_)

# Finaly set the items
b[y, x] = 1

它甚至适用于具有可变大小子列表的indices列表，例如indices = [[1, 3], [0, 1], [0, 2, 3]] 。 如果您的indices列表在每个子列表中始终包含相同数量的项目，那么您可以使用（更有效）以下代码：

b = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])
indices = np.array(indices)
n, m = indices.shape
y = np.arange(n).repeat(m)
x = indices.ravel()
b[y, x] = 1

Answer 3

基于Jérôme 的回答的简单单行代码（要求所有indices项的长度相同）：

>>> b[np.arange(np.size(indices)) // len(indices[0]), np.ravel(indices)] = 1
>>> b
array([[0, 1, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 0]])