Fastest way to filter lists of lists based on a third list?

numpy, performance, python

Solution

All approaches presented here are based on numpys boolean indexing. The approach is to identify matches (independant of row) and then use a reduction (`np.any` or `np.all`) along the rows to see which rows should be eliminated and which should be kept. Finally this mask is applied to your array `A` to get only the valid rows. The only real difference between the approaches is how you create the mask.

Approach 1:

If the values of `B` are known in advance you generally use `|` (or operator) chained comparisons.

a[~np.any(((a == 1) | (a == 2) | (a == 5)), axis=1)]

I'll go through this step-by-step:

Finding matches

>>> ((a == 1) | (a == 2) | (a == 5))
array([[ True,  True],
       [ True, False],
       [False, False],
       [False,  True],
       [False, False]], dtype=bool)

Check each row for one `True`:

>>> np.any(((a == 1) | (a == 2) | (a == 5)), axis=1)
array([ True,  True, False,  True, False], dtype=bool)

Invert it:

>>> ~np.any(((a == 1) | (a == 2) | (a == 5)), axis=1)
array([False, False,  True, False,  True], dtype=bool)

Use boolean indexing:

>>> a[~np.any(((a == 1) | (a == 2) | (a == 5)), axis=1)]
array([[3, 4],
       [6, 7]])

Approach 2:

Instead of these `a == 1 | a == 2 | ...` you could also use `np.in1d`:

>>> np.in1d(a, [1, 2, 5]).reshape(a.shape)
array([[ True,  True],
       [ True, False],
       [False, False],
       [False,  True],
       [False, False]], dtype=bool)

and then use essentially the same approach as above

>>> a[~np.any(np.in1d(a, [1, 2, 5]).reshape(a.shape), axis=1)]
array([[3, 4],
       [6, 7]])

Approach 3:

In case `b` is sorted you can also use `np.searchsorted` to create the mask:

>>> np.searchsorted([1, 2, 5], a, side='left') == np.searchsorted([1, 2, 5], a, side='right')
array([[False, False],
       [False,  True],
       [ True,  True],
       [ True, False],
       [ True,  True]], dtype=bool)

This time you'd need to check if `all` values in reach row are `True`:

>>> b = [1, 2, 5]
>>> a[np.all(np.searchsorted(b, a, side='left') == np.searchsorted(b, a, side='right'), axis=1)]
array([[3, 4],
       [6, 7]])

Timings:

The first approach isn't exactly suitable for arbitary `B` so I don't include it in these timings.

import numpy as np

def setapproach(A, B):  # author: Max Chrétien
    B = set(B)
    indices_to_del = [i for i, sublist in enumerate(A) if B & set(sublist)]
    C = np.delete(A, indices_to_del, 0)
    return C

def setapproach2(A, B):  # author: Max Chrétien & Ev. Kounis
    B = set(B)
    return np.array([sublist for sublist in A if not B & set(sublist)])

def isinapproach(a, b):
    return a[~np.any(np.in1d(a, b).reshape(a.shape), axis=1)]

def searchsortedapproach(a, b):
    b.sort()
    return a[np.all(np.searchsorted(b, a, side='left') == np.searchsorted(b, a, side='right'), axis=1)]

A = np.random.randint(0, 10000, (100000, 2))
B = np.random.randint(0, 10000, 2000)

%timeit setapproach(A, B)
# 929 ms ± 16.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit setapproach2(A, B)
# 1.04 s ± 13.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit isinapproach(A, B)
# 59.1 ms ± 1.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit searchsortedapproach(A, B)
# 56.1 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

The timings, however depend on the range of values, if `B` is already sorted and the lengths of `A`, `B`. But the numpy approaches seams to be almost 20 times faster than the set-solutions. However the difference is mostly because iteration over numpy-arrays with python loops is really inefficient so I'll convert `A` and `B` to `list`s first:

def setapproach_updated(A, B):
    B = set(B)
    indices_to_del = [i for i, sublist in enumerate(A.tolist()) if B & set(sublist)]
    C = np.delete(A, indices_to_del, 0)
    return C

def setapproach2_updated(A, B):
    B = set(B)
    return np.array([sublist for sublist in A.tolist() if not B & set(sublist)])

That may seem strange but let's redo the timings:

%timeit setapproach_updated(A, B)
# 300 ms ± 2.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit setapproach2_updated(A, B)
# 378 ms ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

This is much faster than the plain loops, just by converting it with `tolist` first, but still 5+ times slower than the numpy approaches.

So remember: When you have to use Python-based approaches on NumPy arrays check if it is faster to convert it to a list first!

Let's see how that performs on bigger arrays (these are sizes that approximate those mentioned in your question):

A = np.random.randint(0, 10000000, (1500000, 2))
B = np.random.randint(0, 10000000, 50000)

%timeit setapproach_updated(A, B)
# 4.14 s ± 66.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit setapproach2_updated(A, B)
# 6.33 s ± 95.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit isinapproach(A, B)
# 2.39 s ± 102 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit searchsortedapproach(A, B)
# 1.34 s ± 21.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

The differences got smaller and the `searchsorted`-approach definetly wins.

Approach 4:

I'm not finished yet! Let me surprise you with numba, it's not a lightweight package but extremly powerful if it supports the types and functions you need:

import numba as nb

@nb.njit                # the magic is this decorator
def numba_approach(A, B):
    Bset = set(B)
    mask = np.ones(A.shape[0], dtype=nb.bool_)
    for idx in range(A.shape[0]):
        for item in A[idx]:
            if item in Bset:
                mask[idx] = False
                break
    return A[mask]

Let's see how that performs:

A = np.random.randint(0, 10000, (100000, 2))
B = np.random.randint(0, 10000, 2000)

numba_approach(A, B)   # numba needs a warmup run because it's just-in-time compiling

%timeit numba_approach(A, B)
# 6.12 ms ± 145 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# This is 10 times faster than the fastest other approach!

A = np.random.randint(0, 10000000, (1500000, 2))
B = np.random.randint(0, 10000000, 50000)

%timeit numba_approach(A, B)
# 286 ms ± 16.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# This is still 4 times faster than the fastest other approach!

So, you can make it another order of magnitude faster. Numba doesn't support all python/numpy features (and not all of them are faster) but in this case it's enough!

Problem

Original source