Check if elements in numpy array are between columns of a 2-column numpy array

Viewed 104

I have a 1-D numpy array a with arbitrary size N, and I have other numpy array b with exactly 2 columns and M rows. I want to build a 2-D array with shape (M, N) where the (i,j) element is the bool value b[:,0][i] <= a[j] <= b[:,1][i]. I want a solution that does not go through a for loop.

Example:

import numpy as np

a = np.array([10,  2,  5])

b = np.array([[ 0,  5],[ 5, 10]])


# Expected array

np.array([[False, True, True], [True, False, True]])

1 Answers

The most direct way is to add an axis and broadcast:

(b[:,0][:, np.newaxis] <= a) & (a <= b[:,1][:, np.newaxis])

You can think of the broadcasting in terms of a nested for-loop. You want an output with shape (M,N), you have a with shape (N,) (for j in range(N):) and b[:,0]/b[:,1] with shape (M,) (for i in range(M):). So you would want to cast b[:,0]/b[:,1] to shape (M,1) for another level to the for-loop.

M x 1  for i in range(M):
    N      for j in range(N):
-----
M x N          answer[i,j] = ...
Related