Set 2D array elements based on 2D index array

Question

In Python, say you have a 2D array of zeros of shape (N,4,4):

z = array([[[0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.]],

   [[0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.]]])

and you have a 2D array of indices:

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

and some valued 2D array:

v = array([[ 2.,  4., 10.,  7.],
   [10.,  9.,  9.,  2.],
   [ 3.,  8.,  8.,  8.],
   [ 8.,  6., 10.,  1.]])

Is there a way to fill the elements of z with the values of v but in the slices denoted by i , without using loops?

Note: Is there a way to do this in a scalable fashion such that if you had an N channel array z , where N>>1, you would not need to directly index z[i] when filling it with values from v ?

For clarity, the resulting z array would look like the following:

z = array([[[0., 0., 0., 0.],
    [0., 9., 9., 2.],
    [3., 8., 0., 0.],
    [0., 0., 0., 1.]],

   [[2., 4., 10., 7.],
    [10., 0., 0., 0.],
    [0., 0., 8., 8.],
    [8, 6., 10., 0.]]])

Obviously, using something like:

z = v[i,:,:]

would never work but maybe there's a way to use rows, cols to do this?

Many thanks in advance!

---

Edit:

For clarity here is a similar example but for a 3D z as requested in the comments:

z = array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])


i = array([[2, 1, 2, 1],
       [1, 1, 2, 1],
       [1, 1, 1, 1],
       [1, 0, 0, 1]])


v = array([[5., 5., 0., 4.],
       [4., 6., 8., 3.],
       [4., 0., 4., 8.],
       [7., 6., 5., 7.]])

z would become:

z = array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 6., 5., 0.]],

       [[0., 5., 0., 4.],
        [4., 6., 0., 3.],
        [4., 0., 4., 8.],
        [7., 0., 0., 7.]],

       [[5., 0., 0., 0.],
        [0., 0., 8., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])

Answer 1

You can always use masks and multiply:

from numpy import array
z = array([[[0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.]],

   [[0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.]]])
i = array([[1, 1, 1, 1],
       [1, 0, 0, 0],
       [0, 0, 1, 1],
       [1, 1, 1, 0]])
v = array([[ 2.,  4., 10.,  7.],
   [10.,  9.,  9.,  2.],
   [ 3.,  8.,  8.,  8.],
   [ 8.,  6., 10.,  1.]])
z[0] = ~i.astype(bool)*v
z[1] = i.astype(bool)*v

Output:

z = 
array([[[ 0.,  0.,  0.,  0.],
        [ 0.,  9.,  9.,  2.],
        [ 3.,  8.,  0.,  0.],
        [ 0.,  0.,  0.,  1.]],

       [[ 2.,  4., 10.,  7.],
        [10.,  0.,  0.,  0.],
        [ 0.,  0.,  8.,  8.],
        [ 8.,  6., 10.,  0.]]])

For more general case:

for ax in np.unique(i):
    mask = i == ax
    z[ax] = mask*v
print(z)

Output:

[[[0. 0. 0. 0.]
  [0. 0. 0. 0.]
  [0. 0. 0. 0.]
  [0. 6. 5. 0.]]

 [[0. 5. 0. 4.]
  [4. 6. 0. 3.]
  [4. 0. 4. 8.]
  [7. 0. 0. 7.]]

 [[5. 0. 0. 0.]
  [0. 0. 8. 0.]
  [0. 0. 0. 0.]
  [0. 0. 0. 0.]]]

Note: Here I assumed the values in i are same as there are channels in z . If different, instead of np.unique, you can try:

for i_val,ax in zip(np.unique(i),range(z.shape[0]))::
    mask = i == i_val
    z[ax] = mask*v

EDIT

A one-liner would be:

>>> z = np.array([*map(i.__eq__,np.unique(i))])*np.array([v]*z.shape[0])
>>> z
array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 6., 5., 0.]],

       [[0., 5., 0., 4.],
        [4., 6., 0., 3.],
        [4., 0., 4., 8.],
        [7., 0., 0., 7.]],

       [[5., 0., 0., 0.],
        [0., 0., 8., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])

Answer 2

Possible using numpy.where :

z[0] = np.where(i, z[0], v)
z[1] = np.where(1 - i, z[1], v)

Output:

array([[[ 0.,  0.,  0.,  0.],
        [ 0.,  9.,  9.,  2.],
        [ 3.,  8.,  0.,  0.],
        [ 0.,  0.,  0.,  1.]],

       [[ 2.,  4., 10.,  7.],
        [10.,  0.,  0.,  0.],
        [ 0.,  0.,  8.,  8.],
        [ 8.,  6., 10.,  0.]]])

And after your update, if i hold the indices of z to be changed (and opposite from the 0's and 1's in your example), it can be scaled easily to:

for N in range(z.shape[0]):    
    z[N] = np.where(i == N, v, z[N])

Or if you really want to avoid loops alltogheter:

z = np.where(i == (np.where(np.ones(len(z.reshape(-1))) == 1)[0] // (z.shape[1] * z.shape[2])).reshape(z.shape), v, z)

Ouput:

array([[[0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.],
    [0., 6., 5., 0.]],

   [[0., 5., 0., 4.],
    [4., 6., 0., 3.],
    [4., 0., 4., 8.],
    [7., 0., 0., 7.]],

   [[5., 0., 0., 0.],
    [0., 0., 8., 0.],
    [0., 0., 0., 0.],
    [0., 0., 0., 0.]]])

Answer 3

Something more complicated without numpy:

r = list(zip(i,v))
z[0] = [[l2[i] if l1[i]==0 else 0 for i in range(4)] for l1, l2 in r ]
z[1] = [[l2[i] if l1[i]==1 else 0 for i in range(4)] for l1, l2 in r ]