CODEWITHSUNDEEP
pythonpython-3.xnumpynumpy-ndarray
H123321
H123321

Reputation: 148

Get indices of multidimensional numpy array of arbitrary shape and use it on another array

I have two multidimensional arrays a and b of same shape:

print(a.shape) -> (100, 20, 3, 3)
print(b.shape) -> (100, 20, 3, 3)

Now I only want to get values out of a that match a certain condition, for example:

wanted_value = 5
result_a = a[a[:, :, :, 0] == wanted_value]

That works well, however, now suppose I would like to get data out of another array b of indices that correspond to the indices of found elements in array a. Is it possible to somehow use the indices I can get from searching array a for a certain condition on array b to get data out of array b corresponding to the same indices of the same shape that array a has? I need this because I can not search array b for the condition I need but the data/elements of array b correspond to array a respectively.

Is this possible? I tried something like this:

indices = np.where(a[a[:, :, :, 0] == wanted_value])
results_b = b[indices]

Upvotes: 0

Views: 322

Answers (2)

mujjiga
mujjiga

Reputation: 16876

Numpy supports boolean indexing i.e when you index a numpy array with a Boolean array all the values with truthvalue of True will be returned. This advanced indexing occurs when obj is an array object of Boolean type. The comparison operators return boolean objects and so these can be used to index.

a = np.random.randint(0, 500, (100,20,3,3))
b = np.copy(a)
assert np.array_equal(a[a[:,:,:,0] == 1], b[a[:,:,:,0] == 1])

Check official docs from here.

Upvotes: 1

Karl Knechtel
Karl Knechtel

Reputation: 61526

Yes. It is even simpler than you expect:

b[a[:, :, :, 0] == wanted_value]

The way it works is that a[:, :, :, 0] == wanted_value by itself creates a Numpy array with dtype=bool, referred to in the documentation as a mask. You can apply it to any array of the same shape.

Upvotes: 1

Related Questions