Create numpy array with another array and an indexing array

Question

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?

Why don’t you want to use for loops? You could use generators which use for loops but are more efficient if that’s what you’re looking for. — Asher Mancinelli
– Asher Mancinelli, Commented Nov 9, 2017 at 23:49
I'm just trying to familiarize myself a little more with NumPy. I certainly could use for loops, I just want to know if there's a simple way to write this using NumPy's functions that I'm unaware of. — j_e
– j_e, Commented Nov 9, 2017 at 23:53

Divakar · Accepted Answer · 2017-11-10 03:58:06Z

1

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

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

answered Nov 10, 2017 at 3:58

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Answer 1