In an array of counters that reset, find the start-end index for counter

Question

Given an array that looks like this:

values [0,  0,  1,  2,  3,  4,  5,  0,  1,  2,  3,  4,  5,  6,  0,  0,  1,  2,  3]  
index   0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15  16  17  18

If I search for index 3, I want to get the indexes of the start and end indexes for that counter, before it is reset again, which is 2 - 6.

And for index 10, I want to get 8 - 13. And for index 16, I want to get 16 - 18.

How can I achieve this in numpy?

Answer 1

For example, if you care about index==10, you can try

grp = np.cumsum(np.append(1,np.diff(values)<0))
np.argwhere(grp ==grp[index==10])[1:].T

which gives

array([[ 8,  9, 10, 11, 12, 13]])

Data

values = np.array([0,  0,  1,  2,  3,  4,  5,  0,  1,  2,  3,  4,  5,  6,  0,  0,  1,  2,  3])
index = np.arange(values.size)

Answer 2

You can get the start/end coordinates of the non-null stretches with something like:

idx = np.nonzero(values == 0)[0]
start = idx+1
end = np.r_[idx[1:]-1, len(values)-1]
m = start<end
indices = np.c_[start, end][m]

indices:

array([[ 2,  6],
       [ 8, 13],
       [16, 18]])

Then get the position with searchsorted (assuming you only pass non-zeros indices, else you need an additional check (e.g. is values[position] != 0) and explanation of what should be the output):

indices[np.searchsorted(indices[:, 1],  2)] # [ 2,  6]
indices[np.searchsorted(indices[:, 1], 10)] # [ 8, 13]
indices[np.searchsorted(indices[:, 1], 16)] # [16, 18]

And you can get multiple targets at once:

target = [2, 6, 10, 16]
indices[np.searchsorted(indices[:, 1], target)]

array([[ 2,  6],
       [ 2,  6],
       [ 8, 13],
       [16, 18]])

And if you have indices of zero-values you could mask them in the output:

target = [1, 2, 6, 7, 10, 16]

out = np.ma.masked_array(indices[np.searchsorted(indices[:, 1], target)],
                         np.broadcast_to(values[target, None]==0, (len(target), 2))
                        )

[[-- --]
 [ 2  6]
 [ 2  6]
 [-- --]
 [ 8 13]
 [16 18]]

Used input:

values = np.array([0,  0,  1,  2,  3,  4,  5,  0,  1,  2,  3,  4,  5,  6,  0,  0,  1,  2,  3])