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

In numpy you can set the indices of a 1d array to a value

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

I'm trying to do this with multi-rows and an index for each row in the most programmatically elegant and computationally efficient way possible. For example

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

The desired result is

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

I tried b[indices] and b[:,indices] but they resulted in an error or undesired result.

From searching, there are a few work arounds, but each tends to need at least 1 loop in python.

Solution 1: Run a loop through each row of the 2d array. The draw back for this is that the loop runs in python, and this part won't take advantage of numpy's c processing.

Solution 2: Use numpy put. The draw back is put works on a flattened version of the input array, so the indices need to be flattened too, and altered by the row size and number of rows, which would use a double for loop in python.

Solution 3: put_along_axis seems to only be able to set 1 value per row, so I would need to repeat this function for the number of values per row.

What would be the most computationally and programatically elegant solution? Anything where numpy would handle all the operations?

CodePudding user response：

In [330]: b = np.zeros((3,5),int)

To set the (3,2) columns, the row indices need to be (3,1) shape (matching by broadcasting):

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 does the same thing:

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

CodePudding user response：

On solution to build the indices in each dimension and then use a basic indexing:

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

It works even for indices lists with variable-sized sub-lists like indices = [[1, 3], [0, 1], [0, 2, 3]]. If your indices list always contains the same number of items in each sub-list then you can use the (more efficient) following code:

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

CodePudding user response：

Simple one-liner based on Jérôme's answer (requires all items of indices to be 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]])