CODEWITHSUNDEEP
pythonarraysnumpyvectorization
j_e
j_e

Reputation: 57

Create numpy array with another array and an indexing array

I have an array y that is of shape (m, n) and an array indices that is of shape (m, n, k). Suppose I want to do the following (in Python):

x = np.ndarray(shape=(m,n,k))
for i in range(m):
    for j in range(n):
        for l in range(k):
            x[i,j,l] = y[indices[i,j,l],j]

Is there a way to do this simply using NumPy that doesn't require using for loops?

Upvotes: 1

Views: 550

Answers (1)

Divakar
Divakar

Reputation: 221524

Use NumPy's advanced-indexing for a vectorized assignment -

x = y[indices, np.arange(n)[:,None]]

Upvotes: 1

Related Questions