给定每行的一维索引，如何在二维 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：运行循环遍历二维数组的每一行。这样做的缺点是循环运行在 python，这部分不会利用 numpy 的 c 处理。

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

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

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

Answer 1

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

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 具有 variable-sized sub-lists 的列表，例如 indices = [[1, 3], [0, 1], [0, 2, 3]]。如果您的 indices 列表在每个 sub-list 中始终包含相同数量的项目，那么您可以使用（更有效的）以下代码：

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 2

基于的简单one-liner（要求indices的所有项目为equal-length）：

>>> 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]])

Answer 3

在 [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]])

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

How to set values in a 2d numpy array given 1D indices for each row?

python

arrays

numpy