CODEWITHSUNDEEP
pythonperformancenumpyloops
Learning from masters
Learning from masters

Reputation: 2802

Is it possible to avoid this loop using numpy?

Suppose you have two numpy matrices:

import numpy as np

n = 800
m = 16
A = np.zeros((n, 4)) #np.random.rand(n,4) if you wish
B = np.zeros(m)      #np.random.rand(m) if you wish

In my code, I need to substract each value found in B to all values of A, then calculate that square and get the index that give me the minimum of that new matrix. Therefore, I will have m minimums. From those, I need the minimum. To do this now I use a loop such:

min_C = np.zeros(m)

for j in range(m):
    C = A[:,2]-B[j]
    C = C*C
    min_C[j] = np.where(C == np.amin(C))[0][0]

So I have a minimum for each j. Then, I do the minimum of min_C. Can I avoid this loop to obtain the absolute minimum? I cannot imagine now how to code it, if possible. It's just to improve efficency in my code, which I need it, since this calculation is done lots of times with differents A and B, of course.

Upvotes: 1

Views: 63

Answers (2)

David Wierichs
David Wierichs

Reputation: 545

I'll assume that the fixed index 2 that appears in your sample code will remain fixed, which reduces the problem to the case that A is one-dimensional. If I understand the question correctly and looking at the output of your sample, you are looking for an index and not the value appearing in one of the C's itself, such that we can replace the square with the absolute value. First, we create this absolute value of all combinations of A[i,2] and B[j] via the broadcasting tip by @yatu, and then we compute the minimum of the argmins along the second axis (of size n)

c = np.abs(A[:,2]-B[:,None])
answer = np.min(np.argmin(c, axis=1))

I hope I got correctly what you asked for, if the fixed index of A is to be varied, extend the broadcasting:

c = np.abs(A[:,:]-B[:,None,None])
answer = np.min(np.argmin(c, axis=1))

The output will be an int, which differs from your sample code.

Upvotes: 1

yatu
yatu

Reputation: 88305

Leveraging broadcasting, you can reduce this to the following taking the min of the resulting (m,n) array:

((A[:,2]-B[:,None])**2).min()

Upvotes: 1

Related Questions