python - 访问数组中的左侧和右侧索引，其中元素相差 1

Question

我有一个这样的一维数组:

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

使用下面的函数，我可以访问相邻元素之间的绝对差为 1 的索引对。

当前代码:

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

当前输出:

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

我将如何修改此代码，以便对于差值为 1 的每个位置，我不仅获得一对索引，而且还获得转换左侧的两个索引和右侧的两个索引？

要求的输出:

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

Answer 1

忽略我的变量名。我没有试图表现得专业。此外，这可能以“更短”的方式完成。只是想提供一个解决方案。

代码:

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)

输出:

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