Search code examples
pythonnumpyarray-broadcasting

How to compare the first element with others in numpy?

I'm trying to get an array where it shows whether each of the element equals to the first element in every row:

a = np.array([[1,2,1],[3,3,3],[2,0,1]])
b = a[:,0] == a[:]
print(b)

In 1st row should check 1 == 1, 1 == 2, 1 == 1.

In 2nd row should check 3 == 3, 3 == 3, 3 == 3.

In 3rd row should check 2 == 2, 2 == 0, 2 == 1.

However the output it wrong:

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

Should be:

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

Now what am I doing wrong and how to get the correct result?

Solution

  • To brodcast you need to reshape to add an extra dimension:

    a[:,0,None] == a

    Output:

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