如何在給定每行一維索引的二維 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]])