Laizer
Laizer

Reputation: 6140

Given two vectors, get a matrix of boolean values indicating where elements of vectors are equal

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

Answers (1)

Divakar
Divakar

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

Related Questions