Is it possible to avoid this loop using numpy?

Question

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.

yatu · Accepted Answer · 2020-05-28 16:47:33Z

1

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

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

edited May 28, 2020 at 16:47

answered May 28, 2020 at 16:40

yatu

88.6k12 gold badges93 silver badges148 bronze badges

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

2 Comments

Learning from masters Over a year ago

Wow! really useful! :) Can I also obtain the index? I mean, if I do G = (A[:,2]-B[:,None])**2) I obtain a big matrix where each row is C[j] (well, actually, it is like [ [[ row 1]] [[row 2]] ... ]). It would be also nice to have for each G[j] the minimum in that j

yatu Over a year ago

Yes you have argmin to give you the index of the minimum value. So ((A[:,2]-B[:,None])**2).argmin(). Note that these two have an axis parameter, to find the row-wise or col-wise minimums or argmins @learning too

David Wierichs · Accepted Answer · 2020-05-28 17:27:13Z

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.

Thanks! But I already had the solution with @yatu. Just two things: 1) I actually found the axis I was looking for was 2, not 1. 2) Do you know if np.abs is faster than **2?
Sure, no worries! 1) Okay, but did I get your question right, just out of curiousity. 2) I would not have thought so but a brief check looked like **2 actually is faster, thanks!
yes, thanks! about the axis i said, if you put B = np.zeros(m) it was 1 as you said, but then in my testing file I had B = np.zeros(m,1) and I had to put axis=2 :)

Collectives™ on Stack Overflow

Is it possible to avoid this loop using numpy?

2 Answers 2

2 Comments

3 Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

3 Comments

Related