CODEWITHSUNDEEP
pythonnumpy
MikeRand
MikeRand

Reputation: 4838

numpy: return boolean with true if array elements are in a given set

>>> s = {0, 4, 27}
>>> a = numpy.arange(10)
>>> t = some_func(a)
>>> t
[1, 0, 0, 0, 1, 0, 0, 0, 0, 0]

What's the canonical some_func needed to be to make this work?

What I've tried

I've tried vectorizing a lambda function, which works ... it just doesn't feel like the right way to do this.

>>> f = lambda i: i in s
>>> vf = numpy.vectorize(f)
>>> t = numpy.fromfunction(vf, a.shape)
>>> t.astype(int)
array([1, 0, 0, 0, 1, 0, 0, 0, 0, 0])

Upvotes: 0

Views: 138

Answers (2)

askewchan
askewchan

Reputation: 46568

This would be faster:

t = np.zeros(10, int)
sa = np.array(list(s))
t[sa[sa<10]] = 1

Unless of course, a is not np.arange()

Upvotes: 0

YXD
YXD

Reputation: 32521

Use in1d with s as a NumPy array

>>> import numpy as np
>>> s = np.array([0, 4, 27])
>>> a = np.arange(10)
>>> t = np.in1d(a, s)
>>> t
array([ True, False, False, False,  True, False, False, False, False, False], dtype=bool)
>>>

Upvotes: 3

Related Questions