Find indexes of subarray in numpy array

Question

I have two numpy arrays, one larger, one smaller:

a = np.array([[0,1,0],[0,0,1],[0,1,1]])
b = np.array([[0],[1]])

Is there a function that I can use to find the indexes of the larger array where there is an in an instance of the smaller?

Ideal result:

instances[0] = [[2, 0], [2, 1]]
instances[1] = [[1, 1], [1,2]]

Many thanks!

Answer 1

As far as I know there is not fast numpy function that will do this, but you can loop through and check pretty quickly.

def find_instances(a,b):
    instances = []
    for i in range(a.shape[0] - b.shape[0] + 1):
        for j in range(a.shape[1] - b.shape[1] + 1):
            if np.all(a[i:i+b.shape[0], j:j+b.shape[1]] == b):
                instances.append([i,j])
    return instances

Here each instance is the spot in the top left corner of a that matches the top left corner of b. Not quite the output you requested but it's easy enough to get the rest of the indices if you really need them from there. Hope that helps!