如何使用Numpy查找数组中相同值连续组的索引?
用Numpy查找二进制数组中连续元素的组索引
当然可以用Numpy实现这个需求!其实你提到的np.ediff1d就是核心工具,只是需要搭配一点小技巧来处理数组开头或结尾就有连续目标元素的情况,我来一步步给你拆解实现过程:
步骤1:准备示例数据
先把你的示例数组转换成Numpy数组:
import numpy as np y = np.array([1,1,1,0,0,0,0,0,1,1,0,0,0,0,0,0,1,1,1,1])
步骤2:处理边缘情况——扩展数组
为了确保数组开头或结尾的连续组也能被检测到,我们需要在原数组的前后各添加一个与目标元素(比如1)不同的值。这样不管数组首尾是什么,ediff1d都能捕捉到所有的“起始”和“结束”变化点:
target = 1 # 我们要找连续1的组,找0的话改成0即可 # 前后添加与target相反的值:1 - target 等价于 target ^ 1 extended = np.concatenate(([1 - target], y, [1 - target]))
步骤3:计算差分,定位变化点
用np.ediff1d计算扩展数组的相邻元素差:
diff = np.ediff1d(extended)
这里的差分结果中:
- 当值为
1时,说明前一个元素是0,后一个元素是1——对应连续1组的起始位置 - 当值为
-1时,说明前一个元素是1,后一个元素是0——对应连续1组的结束位置的下一位
步骤4:提取并转换索引
根据差分结果提取起始和结束索引,注意要把结束位置调整回原数组的索引:
# 起始索引:差分等于1的位置,对应原数组的起始点 starts = np.where(diff == 1)[0] # 结束索引:差分等于-1的位置减1,得到原数组中连续组的最后一个元素索引 ends = np.where(diff == -1)[0] - 1
步骤5:配对成组索引
把起始和结束索引配对成元组列表,就是我们要的结果:
groups = list(zip(starts, ends)) print(groups) # 输出:[(0, 2), (8, 9), (16, 19)]
封装成通用函数
为了方便复用,我们可以把上面的逻辑封装成一个函数,支持查找连续1或0的组:
def find_consecutive_groups(arr, target): # 扩展数组处理边缘情况 extended = np.concatenate(([1 - target], arr, [1 - target])) diff = np.ediff1d(extended) # 提取起始和结束索引 starts = np.where(diff == 1)[0] ends = np.where(diff == -1)[0] - 1 # 返回配对后的组索引 return list(zip(starts, ends))
测试一下找连续0的情况:
zero_groups = find_consecutive_groups(y, 0) print(zero_groups) # 输出:[(3, 7), (10, 15)]
这个方法完全基于Numpy的向量化操作,效率非常高,就算处理超大数组也比纯Python循环快得多,完美解决你的需求~
内容的提问来源于stack exchange,提问作者tiefenauer
相关产品推荐
相关产品推荐

