Home > Net >  Given 2 numpy 1-D sorted arrays, check for each element in A if there is an element in B such that a
Given 2 numpy 1-D sorted arrays, check for each element in A if there is an element in B such that a

Time:09-21

I have two 1-D numpy sorted arrays, A and B. For each element in A, I want to check if there is an element in B such that

abs(a-b) <= 10

and keep only the elements that match the condition. Is there a way to do it without loops?

The arrays are not necessarily of the same length, and the idx of the element that is closest to A in B could be far from the original idx in A. Otherwise I could use something like

A[(A >= B - 10) & (A <= B   10)]   

I don't mind using a lot of space, so I can create 2 new arrays of size B[-1] and use the elements in A,B as indices where there is 1 in the new arrays, but I can't think of a solution without loops that way either. Thought of using np.lib.stride_tricks.as_strided, but not sure how.

Thanks.

CodePudding user response:

This should do it:

answer = a[np.sum(abs(a.reshape(a.shape[0], 1) - b) <= 10,axis=-1) > 0]

CodePudding user response:

You can use np.isclose with broadcasting B

A[np.isclose(A, B[:, None], atol=10).any(0)]
  • Related