如何获取数组中相邻元素差为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])

现有代码用来提取相邻元素绝对差为1的过渡索引对：

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

当前输出：

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

需要修改代码，让每个过渡位置不仅返回索引对，还要包含过渡位置左侧两个索引和右侧两个索引，预期输出：

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

解决方案代码 #

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])
# 定位所有过渡的起始索引
transition_indices = np.where(np.abs(np.diff(data)) == 1)[0]
# 构建目标结果列表
result = []
for i in transition_indices:
    # 添加左侧两个索引
    result.extend([i-2, i-1])
    # 添加过渡索引对
    result.append((i, i+1))
    # 添加右侧两个索引
    result.extend([i+1, i+2])

print(result)

代码说明 #

1. 先用np.where找到所有满足条件的过渡起始索引i（也就是相邻元素差为1的左边元素索引）；
2. 对每个i，依次往结果列表里添加：左侧两个索引（i-2、i-1）、过渡索引对(i, i+1)、右侧两个索引（i+1、i+2）；
3. 最终合并后的列表就是预期的输出格式。

内容的提问来源于stack exchange，提问作者Bhar Jay