Numpy 2D-Array - Set values with index array

Question

I'm wondering if there is a better way of assigning values from a numpy array to a 2D numpy array based on an index array.

arr1 = np.array([1, 1, 1])
arr2 = np.array([1, 1, 1])
arr3 = np.array([1, 1, 1])

index = np.array([1, 2, 1])
values = np.array([0, 0, 0])

for idx, arr_x in enumerate(newest):
    arr_x[index[idx]] = values[idx]

Right now I'm doing it like this to get an result like:

[[1 0 1]
 [1 1 0]
 [1 0 1]]

I thought newest[:, index] = values should work but it's not.

Does anybody know how to do this? I'm sure there is a better way. Thanks in advance.

Answer 1

Use np.put_along_axis. You need to use np.expand_dims to make all the variable has the same number of dimensions i.e. 2 in your case. Initially index and values have a shape of (3,), which needs to be converted to (3,1).

>>> arr = np.vstack((arr1, arr2, arr3))
>>> index = np.array([1, 2, 1])
>>> index = np.expand_dims(index, axis=1)
>>> values = np.array([0, 0, 0])
>>> values = np.expand_dims(values, axis=1)
>>> np.put_along_axis(arr, index, values, axis=1)
>>> arr
array([[1, 0, 1],
       [1, 1, 0],
       [1, 0, 1]])

Answer 2

Not so much more efficient but slightly better given how your data is formatted

arr1 = np.array([1, 1, 1])
arr2 = np.array([1, 1, 1])
arr3 = np.array([1, 1, 1])

index = np.array([1, 2, 1])

then

for arr,m in zip((arr1,arr2,arr3),index):
    arr[m] = 0

final_array = np.stack((arr1,arr2,arr3))

Answer 3

You need to replace the slice with actual row indices:

In [66]: newest = np.ones((3,3),int)
In [67]: newest[np.arange(3),index] = values
In [68]: newest
Out[68]: 
array([[1, 0, 1],
       [1, 1, 0],
       [1, 0, 1]])

This sets (0,1), (1,2), and (2,1) values.

put_along_axis does this as well, but this is a simple enough case that the direct indexing is easy.