Search code examples
pythonnumpyarray-broadcasting

Convert a simple for loop using numpy's array broadcasting

I have two arrays with shapes similar to those given in the code block below. I need to broadcast the operation in the for loop to improve the performance of my code (this operation is performed millions of times)

I know it looks very simple but I've been trying for the last 30 minutes and I can't find the solution.

import numpy as np

# Input arrays
N1, N2, N3 = 500, 2, 1700
aa = np.random.uniform(0., 1., (N1, N2, N3))
bb = np.random.randint(0, N1, N3)

# Block that needs to be broadcasted to improve the performance
arr = []
for i, x in enumerate(aa.T):
    arr.append(x[:, bb[i]])

# Final array
arr = np.array(arr).T
Solution 
  • aa[bb,:,np.arange(aa.shape[2])].T