Accessing left and right side indices in an array, where the elements differ by 1-CodePudding

I have a one-dimensional array like this:

data=np.array([1,1,1,1,0,0,0,0,0,1,1,1,4,4,4,4,4,4,1,1,1,0,0,0])

With the function below, I am able to access the index pairs where the absolute difference between adjacent elements is 1.

Current code:

result = [(i, i   1) for i in np.where(np.abs(np.diff(data)) == 1)[0]]

Current output:

[(3, 4), (8, 9), (20, 21)]

How would I modify this code so that for each place where the difference is 1, I get not only the pair of indices, but also two indices to the left and two to the right of the transition?

Required output:

[2,3,(3, 4),4,5,7,8,(8, 9),9,10,19,20, (20, 21),21,22]

CodePudding user response：

Ignore my variable names. I didn't try to be professional. Also, this could probably be done in a "shorter" way. Just wanted to provide a solution.

Code:

import numpy as np

data=np.array([1,1,1,1,0,0,0,0,0,1,1,1,4,4,4,4,4,4,1,1,1,0,0,0])

result = [(i - 1, i, (i, i   1), i   1, i   2) for i in np.where(np.abs(np.diff(data)) == 1)[0]]
new_result = []
for r in result:
    for r1 in r:
        new_result.append(r1)
new_result = np.array(new_result, dtype=object)

print(new_result)

Output:

[2, 3, (3, 4), 4, 5, 7, 8, (8, 9), 9, 10, 19, 20, (20, 21), 21, 22]