Python Numpy - 3-dimensional indices in 2-dimensional array without loop

Question

I want to construct an array V1, of shape (n,p,q) using an array of indices, idx, with the same shape, applied to an array V0, of shape (p,q). The way to construct it with a loop is the following.

for i in range(n):
    V1[i,:,:] = V0[idx[i,:,:],range(q)]

In other words, the idx[i,:,:] array contains indices for the elements of the 1st dimension of V0. I apply it with the associated index of the 2nd dimension, captured in range(q), to get the corresponding element along the fist dimension of the final array V1.

My question is the following: is there a way to construct V1 without looping, by using broadcasting/indexing techniques?

Thank you.

score 0 · Answer 1 · answered Oct 21 '18 at 23:52

How about just

V1 = V0[idx, range(q)] #?

Example:

import numpy as np

# set up dummy data
n,p,q = 3,4,5
V1 = np.empty((n,p,q))
V0 = np.random.rand(p,q)
idx = np.random.randint(0,n,(n,p,q))

# original
V1_old = V1.copy()
for i in range(n):
    V1_old[i,:,:] = V0[idx[i,:,:],range(q)]

# new
V1_new = V0[idx, range(q)]

# test
print(np.array_equal(V1_old, V1_new)) # True

Python Numpy - 3-dimensional indices in 2-dimensional array without loop

1 Answers1