Assigning numpy array values based on complex indexing scheme-CodePudding

I am attempting to populate a first [m x n] numpy array with values drawn from a second [m x n x f] array, where the index f is supplied by a third [m x n] array, and I’ve not yet been able to find a way to do this without expensive for loops. Here is a simple representative example:

import numpy as np
import random

m, n, k  = 3, 2, 4
np.random.seed(1234)
a = np.random.randint(0,9,(m, n, k))
print(a)
b = np.random.randint(0,k,(m,n))
print(b)
c = np.zeros((m, n))
print(c)

I would now like to populate c with values from a as follows:

c[i,j] = a[i,j,b[i,j]]

In other words, for each position [i,j ] in c, draw the value from a at the same position [i,j] along the first two axes, and at a position along the third axis given by evaluating the array b at the [i,j] position. In my example case (your random ints may vary?), the arrays are valued as follows:

a :
[[[3 6 5 4]
  [8 1 7 6]]
 [[8 0 5 0]
  [6 2 0 5]]
 [[2 6 3 7]
  [0 0 3 2]]]

b :
[[3 0]
 [1 3]
 [3 3]]

The desired result is:

c :
[[4 8]
 [0 5]
 [7 2]]

We arrive at this result as follows: start with i, j = 0,0. Then, c[0,0] = [a[0,0,b[0,0]], and since b[0,0] is valued at 3, a is evaluated at a[0,0,3], which is 4. Proceed to i, j = 0, 1. Here c[0,1] = [a[0,1,b[0,1]], and since b[0,1] is valued at 0, a is evaluated at a[0,1,0], which is 8. And so on. It is very clear to me how to do this with nested loops, but I am trying to avoid this since my actual arrays are much larger. Many thanks in advance for your help!

CodePudding user response：

Use np.unravel_index to create the indexes, index and then .reshape:

c = np.empty((m, n), dtype=np.int32)
ro, co = np.unravel_index(np.arange(n * m), (m, n))
c[:, :] = a[ro, co, b[ro, co]].reshape(3,2)
print(c)

Output

[[4 8]
 [0 5]
 [7 2]]

CodePudding user response：

Numpy indices broadcast to the shape of the output array:

c = a[np.arange(b.shape[0])[:, None], np.arange(b.shape[1]), b]

You can also use np.take_along_axis to avoid manually generating the index in the other axes, but here the index has to broadcast to the input array:

c = np.take_along_axis(a, b[..., None], -1)