Row-wise replacement of numpy array with values of another numpy array

Question

I have 0 s and 1 s store in a 3-dimensional numpy array:

g = np.array([[[0, 1], [0, 1], [1, 0]], [[0, 0], [1, 0], [1, 1]]])
# array([
#     [[0, 1], [0, 1], [1, 0]],
#     [[0, 0], [1, 0], [1, 1]]])

and I'd like to replace these values by those in another array using a row-wise replacement strategy. For example, replacing the vales of g by x :

x = np.array([[2, 3], [4, 5]])
array([[2, 3],
       [4, 5]])

to obtain:

array([
     [[2, 3], [2, 3], [3, 2]],
     [[4, 4], [5, 4], [5, 5]]])

The idea here would be to have the first row of g replaced by the first elements of x ( 0 becomes 2 and 1 becomes 3 ) and the same for the other row (the first dimension - number of "rows" - will always be the same for g and x )

I can't seem to be able to use np.where because there's a ValueError: operands could not be broadcast together with shapes (2,3,2) (2,2) (2,2) .

Answer 1

IIUC,

np.stack([x[i, g[i]] for i in range(x.shape[0])])

Output:

array([[[2, 3],
        [2, 3],
        [3, 2]],

       [[4, 4],
        [5, 4],
        [5, 5]]])

Answer 2

Vectorized approach with np.take_along_axis to index into the last axis of x with g using axis=-1 -

In [20]: np.take_along_axis(x[:,None],g,axis=-1)
Out[20]: 
array([[[2, 3],
        [2, 3],
        [3, 2]],

       [[4, 4],
        [5, 4],
        [5, 5]]])

Or with manual integer-based indexing -

In [27]: x[np.arange(len(g))[:,None,None],g]
Out[27]: 
array([[[2, 3],
        [2, 3],
        [3, 2]],

       [[4, 4],
        [5, 4],
        [5, 5]]])

Answer 3

From what I understand, g is an array of indexes (indexes being 0 or 1) and x is the array to who's values you use.

Something like this should work (tested quickly)

import numpy as np

def swap_indexes(index_array, array):
   out_array = []
   for i, row in enumerate(index_array):
         out_array.append([array[i,indexes] for indexes in row])
   return np.array(out_array)


index_array = np.array([[[0, 1], [0, 1], [1, 0]], [[0, 0], [1, 0], [1, 1]]])
x = np.array([[2, 3], [4, 5]])
print(swap_indexes(index_array, x))

[EDIT: fixed typo that created duplicates]

Answer 4

One solution, is to simply use comprehension directly here:

>>> np.array([[x[i][c] for c in r] for i, r in enumerate(g)])
array([[[2, 3],
        [2, 3],
        [3, 2]],

       [[4, 4],
        [5, 4],
        [5, 5]]])