numpy find slice along an axes where the first and last occurring value occurs

user_2713740 0
Viewed 63

I have a 3D numpy array with integer values, something defined as:

import numpy as np
x = np.random.randint(0,  100, (10, 10, 10))

Now what I want to do is find the last slice (or alternatively the first slice) along a given axes (say 1) where a particular value occurs. At the moment, I do something like:

first=None
last=None
val = 20

for i in range(len(x.shape[1]):
    slice = x[:, i, :]
    if len(slice[slice==val]) > 0:
        if not first:
            first = i
        last = i

return first, last

This seems a bit unpythonic and I wonder if there is some numpy magic to do this?

numpy python
2 Answers

You probably can optimize this to be faster, but here is a vectorized version of what you search: 

axis = 1
mask = np.where(x==val)[axis]
first, last = np.amin(mask), np.amax(mask)

It first finds the element val in your array using np.where and returns the min and max of indices along desired axis.

user_49759810

Per your question, you want to check if there's any such valid slice and hence, get the start/first, stop/last indices. In absence of any such valid slice, we must return None's there. That needs extra check. Also, we can use masking to get those indices in an efficient manner, like so -

def slice_info(x, val):
    n = (x==val).any((0,2))
    if n.any():
        return n.argmax(), len(n)-n[::-1].argmax()-1
    else:
        return None,None

Benchmarking

Other proposed solution(s) :

# https://stackoverflow.com/a/61629916/ @Ehsan
def where_amin_amax(x, val):
    axis = 1
    mask = np.where(x==val)[axis]
    first, last = np.amin(mask), np.amax(mask)
    return first, last

Timings -

# Same setup as in given sample
In [157]: np.random.seed(0)
     ...: x = np.random.randint(0,  100, (10, 10, 10))

In [158]: %timeit where_amin_amax(x, val=20)
     ...: %timeit slice_info(x, val=20)
15.1 µs ± 287 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
9.63 µs ± 43.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# Bigger
In [159]: np.random.seed(0)
     ...: x = np.random.randint(0,  100, (100, 100, 100))

In [160]: %timeit where_amin_amax(x, val=20)
     ...: %timeit slice_info(x, val=20)
3.34 ms ± 31.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
691 µs ± 3.69 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
user_32938810
Related