find numpy array in other numpy array

Question

I need to find a small numpy array in a much larger numpy array. For example:

import numpy as np
a = np.array([1, 1])
b = np.array([2, 3, 3, 1, 1, 1, 8, 3, 1, 6, 0, 1, 1, 3, 4])

A function

find_numpy_array_in_other_numpy_array(a, b)

should return indices

[3, 4, 11]

that represent where the complete numpy array a appears in the complete numpy array b.

There is a brute force approach to this problem that is slow when dealing with very large b arrays:

ok = []
for idx in range(b.size - a.size + 1):
    if np.all(a == b[idx : idx + a.size]):
        ok.append(idx)

I am looking for a much faster way to find all indices of the full array a in array b. The fast approach should also allow other comparison functions, e.g. to find the worst case difference between a and b:

diffs = []
for idx in range(b.size - a.size + 1):
    bi = b[idx : idx + a.size]
    diff = np.nanmax(np.abs(bi - a))
    diffs.append(diff)

Answer 1

Generic solution setup

For a generic solution, we can create 2D array of sliding windows and then perform the relevant operations -

from skimage.util.shape import view_as_windows

b2D = view_as_windows(b,len(a))

NumPy equivalent implementation.

Problem #1

Then, to solve for matching indices problem, it's simply -

matching_indices = np.flatnonzero((b2D==a).all(axis=1))

Problem #2

To solve for the second problem, it maps easily by keeping in mind that any ufunc reduction operation to get an output element is to be translated into reduction along the second axis in the proposed solution using that ufunc's axis argument -

diffs = np.nanmax(np.abs(b2D-a),axis=1)