How to do indexing of a NumPy 3D-array based on 2D-array in Python?

user_10069252 0
Viewed 710

Let say I have a NumPy array A of shape (66,5) and B of shape (100, 66, 5).

The elements of A will index the first dimension (axis=0) of B, where the values are from 0 to 99 (i.e. the first dimension of B is 100).

A = 
array([[   1,    0,    0,    1,    0],
       [   0,    2,    0,    2,    4],
       [   1,    7,    0,    5,    5],
       [   2,    1,    0,    1,    7],
       [   0,    7,    0,    1,    4],
       [   0,    0,    3,    6,    0]
       ....                         ]])

For example, A[4,1] will take index 7 of the first dimension of B, index 4 of the second dimension of B and index 1 of the third dimension B.

What I wanted to is to produce array C of shape (66,5) where it contains the elements in B that are selected based on the elements in A.

numpy arrays python
3 Answers

You can use np.take_along_axis to do that:

import numpy as np
np.random.seed(0)
a = np.random.randint(100, size=(66, 5))
b = np.random.random(size=(100, 66, 5))
c = np.take_along_axis(b, a[np.newaxis], axis=0)[0]
# Test some element
print(c[25, 3] == b[a[25, 3], 25, 3])
# True
user_17827920

If I understand correctly, you are looking for advances indexing of first dimension of B. You can use np.indices to create the indices required for the other two dimensions of B and use advanced indexing:

idx = np.indices(A.shape)
C = B[A,idx[0],idx[1]]

Example:

B = np.random.rand(10,20,30)
A = np.array([[   1,    0,    0,    1,    0],
       [   0,    2,    0,    2,    4],
       [   1,    7,    0,    5,    5],
       [   2,    1,    0,    1,    7],
       [   0,    7,    0,    1,    4],
       [   0,    0,    3,    6,    0]])

print(C[4,1]==B[7,4,1])
#True
user_49759810

Use the following (using functions of NumPy library):

print(A)
# array([[2, 0],
#        [1, 1],
#        [2, 0]])

print(B)
# array([[[ 5,  7],
#         [ 0,  0],
#         [ 0,  0]],

#        [[ 1,  8],
#         [ 1,  9],
#         [10,  1]],

#        [[12, 22],
#         [ 2,  2],
#         [ 2,  2]]])

temp = A.reshape(-1) + np.cumsum(np.ones([A.reshape(-1).shape[0]])*B.shape[0], dtype = 'int') - 3
C = B.swapaxes(0, 1).swapaxes(2, 1).reshape(-1)[temp].reshape(A.shape)

print(C)
# array([[12,  7],
#        [ 1,  9],
#        [ 2,  0]])
user_137205650
Related