find which samples are in an integer pool

user_353337 0
Viewed 32

I have two arrays of integers, one very large pool (can have hundreds of thousands of entries) and one samples (can also have many, but only in the order of hundreds/thousands). For each entry in samples, I need to find out if it is in pool. I can do an explicit loop:

import numpy as np


rng = np.random.default_rng(0)

n = 1000
m = 10
pool = rng.integers(0, 100000, n)
samples = rng.integers(0, 100000, m)

is_in_pool = [s in pool for s in samples]

There certainly is a faster variant though.

Any hints?

numpy python
2 Answers

The numpy way to do it, is to use isin:

np.isin(samples, pool)

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

However this is not faster than the list comprehension, which is optimized in C, too. A little speedup can be gained if both arrays are unique using assume_unique=True in np.isin.
The fastest I could get is to convert the pool into a set which as a hash map has constant lookup time. The conversion to a set takes some time and is not included in the following performance measurement of the different approaches which makes the approach using a set not a fair comparison.

pool_set = set(pool)

%%timeit
[s in pool_set for s in samples]
1.81 µs ± 8.17 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)


%%timeit
[s in pool for s in samples]
44.5 µs ± 3.28 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)


%%timeit
np.isin(samples, pool, assume_unique=True)
57.7 µs ± 922 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)


%%timeit
np.isin(samples, pool)
87.9 µs ± 3.09 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
user_98592130

As a hint, you can speed this calculation up by removing duplicates from the arrays.

Hint 2: There is a numpy function that can answer this for you.

user_37809350
Related