CODEWITHSUNDEEP
numpy
Arnaud
Arnaud

Reputation: 467

How to check the presence of a given numpy array in a larger-shape numpy array?

I guess the title of my question might not be very clear..

I have a small array, say a = ([[0,0,0],[0,0,1],[0,1,1]]). Then I have a bigger array of a higher dimension, say b = ([[[2,2,2],[2,0,1],[2,1,1]],[[0,0,0],[3,3,1],[3,1,1]],[...]]).

I'd like to check if one of the elements of a can be found in b. In this case, I'd find that the first element of a [0,0,0] is indeed in b, and then I'd like to retrieve the corresponding index in b.

I'd like to do that avoiding looping, since from the very little I understood from numpy arrays, they are not meant to be iterated over in a classic way. In other words, I need it to be very fast, because my actual arrays are quite big.

Any idea? Thanks a lot!

Arnaud.

Upvotes: 1

Views: 62

Answers (1)

John Kealy
John Kealy

Reputation: 1883

I don't know of a direct way, but I here's a function that works around the problem:

import numpy as np

def find_indices(val, arr):
    # first take a mean at the lowest level of each array,
    # then compare these to eliminate the majority of entries
    mb = np.mean(arr, axis=2); ma = np.mean(val)
    Y = np.argwhere(mb==ma)
    indices = []
    # Then run a quick loop on the remaining elements to
    # eliminate arrays that don't match the order
    for i in range(len(Y)):
        idx = (Y[i,0],Y[i,1])
        if np.array_equal(val, arr[idx]):
            indices.append(idx)

    return indices

# Sample arrays
a = np.array([[0,0,0],[0,0,1],[0,1,1]])
b = np.array([ [[6,5,4],[0,0,1],[2,3,3]],   \
               [[2,5,4],[6,5,4],[0,0,0]],   \
               [[2,0,2],[3,5,4],[5,4,6]],   \
               [[6,5,4],[0,0,0],[2,5,3]]  ])



print(find_indices(a[0], b))
# [(1, 2), (3, 1)]
print(find_indices(a[1], b))
# [(0, 1)]

The idea is to use the mean of each array and compare this with the mean of the input. np.argwhere() is the key here. That way you remove most of the unwanted matches, but I did need to use a loop on the remainder to avoid the unsorted matches (this shouldn't be too memory-consuming). You'll probably want to customise it further, but I hope this helps.

Upvotes: 1

Related Questions