Search code examples
pythonnumpymatrixnested-loops

Avoiding nested for loops in Python

I have a matrix A with dimension of 1024 * 307200 and another matrix B of dimension 1024 * 50. I am performing L2_norm on these two matrices in a nested for loop to get my final matrix C as 307200 * 50.
You can find the code below:

    for i in range(307200):
        for l in range(50):
            C[i,l] =  numpy.linalg.norm(A[:,i] - B[:,l]))

As you see the dimension of my variables is huge which is leading to a very high latency. I want to avoid this nested loop since for each values of i and l, I am using the same function.

Is there any way to optimize the above loop?

Solution

  • Maybe you could replace the inner loop and your function with these matrix operations?

    for i in range(307200):
    temp = A[:,i,np.newaxis] - B[:]
    C[i,:] = np.linalg.norm(temp, axis=0)

    For smaller arrays, I got about 20 times lower running times. Perhaps you gain even more. In any case, plese make sure that you are receiving good results (on a smaller sets).