How to compare an array against a list of arrays?

Viewed 165

Suppose I have a list that contains a bunch of numpy ndarrays (or even torch Tensors):

a, b, c = np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3)
collection = [a, b, c]

Now if I was to check whether array b was in collection (assuming I don't know what arrays exist in collection), then attempting: b in collection spits out the following error:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

And the same would apply to tuples containing arrays.

One way around this was to do a list comprehension:

True in [(b == x).all() for x in collection]

However this requires a for loop and I was wondering if there was a more "efficient" way to accomplish this?

3 Answers

I would use numpy arrays the whole way through:

import numpy as np
a, b, c = np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3)
array = np.dstack([a, b, c])
print(array.shape)
# (3, 3, 3)
np.all(array == b, axis=1).all(axis=1).any()
# True

You can stack arbitrary shaped tensors along axis=0 in a numpy array and then compare all remaining axes at once with np.all (this is just a slightly clearer version of PaulH's answer):

ugly_shaped_tensor_list = [np.random.rand(3,7,5,3) for j in range(10)]
known_tensor = ugly_shaped_tensor_list[1]

# stack all tensors in a single array along axis=0:
tensor_stack = np.stack(ugly_shaped_tensor_list)

# compare all axes except the "list" axis, 0:
matches = np.all(tensor_stack == known_tensor, axis=(1,2,3,4))
# array([False,  True, False, False, False, False, False, False, False, False])
matches.any()
# True

Okay this was much simpler than anticipated...

You can literally just stack your arrays/tensors into a higher dimension (depth/channels in this case) and then the resultant is an array that is a collection of all the other arrays but stored independently in a "different dimension".

a, b, c = np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3)
collection = np.stack((a, b, c))

Now you can simply check for b in collection as if you were comparing it against a list:

b in collection
> True
Related