How do I search for index values in each row of an array in numpy creating a boolean array

Question

Given an array with size MxN and an array with size Mx1, I want to compute a boolean array with MxN.

import numpy as np

M = 2
N = 3

a = np.random.rand(M, N) # The values doesn't matter
b = np.random.choice(a=N, size=(M, 1), replace=True)
# b =
# array([[2],
#        [1]])

# I found this way to compute the boolean array but I wonder if there's a fancier, elegant way

index_array = np.array([np.array(range(N)), ]*M)
# Create an index array
# index_array = 
# array([[0, 1, 2],
#        [0, 1, 2]])
#

boolean_array = index_array == b
# boolean_array =
# array([[False, False,  True],
#        [False,  True, False]])
#

So i wonder if theres's a fancier, pythonic way of doing this

Answer 1

You could simplify by leveraging broadcasting an comparing with a single 1d range directly:

M = 2
N = 3
a = np.random.rand(M, N) 
b = np.random.choice(a=N, size=(M, 1), replace=True)

print(b)
array([[1],
       [2]])

b == np.arange(N)
array([[False,  True, False],
       [False, False,  True]])

---

In general, broadcasting is handy in these cases because it saves us from having to create arrays compatible in shape to perform operations with other arrays. For the generated array, I'd probably go with the following instead:

np.broadcast_to(np.arange(N), (M,N))
array([[0, 1, 2],
       [0, 1, 2]])

Though as mentioned, NumPy makes life easier here so that we don't have to worry about that.