Creating a new matrix from a matrix of index in numpy

Question

I have a 3D numpy array A with shape(k, l, m) and a 2D numpy array B with shape (k,l) with the indexes (between 0 and m-1) of particular items that I want to create a new 2D array C with shape (k,l), like this:

import numpy as np
A = np.random.random((2,3,4))
B = np.array([[0,0,0],[2,2,2]))
C = np.zeros((2,3))
for i in range(2):
    for j in range(3):
        C[i,j] = A[i, j, B[i,j]]

Is there a more efficient way of doing this?

Answer 1

Use inbuilt routine name fromfunction of Numpy library. And turn your code into

C = np.fromfunction(lambda i, j: A[i, j, B[i,j]], (5, 5))

Answer 2

Setup:

import numpy as np
k,l,m = 2,3,4
a = np.arange(k*l*m).reshape(k,l,m)
b = np.random.randint(0,4,(k,l))

print(a)
print('*'*10)
print(b)

---

[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]
**********
[[3 0 3]
 [2 1 2]]

---

Use integer indexing to select the values then reshape.

x,y = np.indices(a.shape[:-1])
c = a[x,y,b]
print(c)

---

[[ 3  4 11]
 [14 17 22]]

---

Usingnumpy.ix_ .

x,y = np.ix_(np.arange(a.shape[0]),np.arange(a.shape[1]))
d = a[x,y,b]