Jack Rolph
Jack Rolph

Reputation: 595

Find indexes of subarray in numpy array

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!

Upvotes: 4

Views: 1528

Answers (1)

Erik Parkinson
Erik Parkinson

Reputation: 356

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!

Upvotes: 2

Related Questions