Reputation: 6140
I have two vectors: A
of length m
, and B
of length n
.
I would like to get an m
by n
matrix whose element (x,y)
is equal to A[x] == B[y]
.
What's the fastest way to construct this matrix in numpy?
Upvotes: 3
Views: 363
Reputation: 221524
Use NumPy's broadcasting feature
by extending the A
to a 2D
version with None/np.newaxis
and then compare against 1D
array B resulting in the desired 2D
boolean array -
A[:,None] == B
An alternative to creating that 2D
version of A
would be with reshaping such that the second axis is of length = 1
. So, an alternative solution would be -
A.reshape(-1,1) == B
Sample run -
In [89]: A
Out[89]: array([0, 1, 2, 3])
In [90]: B
Out[90]: array([4, 3, 2, 1, 0])
In [91]: A[:,None] == B
Out[91]:
array([[False, False, False, False, True],
[False, False, False, True, False],
[False, False, True, False, False],
[False, True, False, False, False]], dtype=bool)
Upvotes: 4