CODEWITHSUNDEEP
pythonnumpy
Greynes
Greynes

Reputation: 749

Apply an offset to the indices obtained from np.where

I have a 3d numpy array and I obtain the indices that meet a certain condition, for example:

a = np.tile([[1,2],[3,4]],(2,2,2))
indices = np.where(a == 2)

To this indices, I need to apply an offset, fo example (0, 0, 1), and view if meet another condition.

Something like this:

offset = [0, 0, 1]
indices_shift = indices + offset

count = 0
for i in indices_shift:
    if a[i] == 3:
       count += 1

In this example, with the offset of (0,0,1), the indices looks like:

indices = (array([0, 0, 0, 0, 1, 1, 1, 1], dtype=int64), array([0, 0, 2, 2, 0, 0, 2, 2], dtype=int64), array([1, 3, 1, 3, 1, 3, 1, 3], dtype=int64))

and I think that adding the offset the results should be something like:

indices_shift = (array([0, 0, 0, 0, 1, 1, 1, 1], dtype=int64), array([0, 0, 2, 2, 0, 0, 2, 2], dtype=int64), array([2, 4, 2, 4, 2, 4, 2, 4], dtype=int64))

Is there any easy way to do that?

Thanks.

Upvotes: 1

Views: 753

Answers (1)

Divakar
Divakar

Reputation: 221614

Here's one approach -

idx = np.argwhere(a == 2)+[0,0,1]
valid_mask = (idx< a.shape).all(1)
valid_idx = idx[valid_mask]
count = np.count_nonzero(a[tuple(valid_idx.T)] == 3)

Steps :

  • Get the indices for matches against 2. Use np.argwhere here to get those in a nice 2D array with each column representing an axis. Another benefit is that this makes it generic to handle arrays with generic number of dimensions. Then, add offset in a broadcasted manner. This is idx.

  • Among the indices in idx, there would be few invalid ones that go beyond the array shape. So, get a valid mask valid_mask and hence valid indices valid_idx among them.

  • Finally index into input array with those, compare against 3 and count the number of matches.

Upvotes: 1

Related Questions