jjepsuomi
jjepsuomi

Reputation: 4373

How to find the set of indices where two vectors have equal elements in Python

I have two vectors in Python: Predictions and Labels. What I would like to do is to find out the set of indices where these two vectors have equal elements. For example, lets say the vectors are:

Predictions = [4, 2, 5, 8, 3, 4, 2, 2]

     Labels = [4, 3, 4, 8, 2, 2, 1, 2]

So the set of indices where the two vectors have equal elements would be:

Indices = [0, 3, 7]

How can I get this in Python? Without using for-loops etc. Is there a built-in function for example in numpy?

Thank you for any help!

Upvotes: 7

Views: 13625

Answers (2)

daniel451
daniel451

Reputation: 10992

For two arrays a, b with:

a = np.array([1, 2, 3, 4, 5])
b = np.array([1, 3, 2, 4, 5])

np.equal(a,b) has the same output as a==b (which is easier to understand at first, I think):

> array([ True, False, False,  True,  True], dtype=bool)

The elements are checked element-wise and then an array of booleans is created.

np.where() checks some condition element-wise on an array:

np.where(a > 2)
> (array([2, 3, 4]),)

So combining np.where and np.equal is the thing you want:

np.where(np.equal(a,b))
> (array([0, 3, 4]),)

edit: nevermind, just saw that I were far too slow ^^

Upvotes: 5

Andrzej Pronobis
Andrzej Pronobis

Reputation: 36086

This is one way of doing it with numpy:

np.where(np.equal(Predictions, Labels))

which is equivalent to:

np.equal(Predictions, Labels).nonzero()

It will return a single element tuple though, so to get the actual array, add [0] as in:

np.equal(Predictions, Labels).nonzero()[0]

Upvotes: 7

Related Questions