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?
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])