判断整数列表中非零元素是否构成单个连续组的优化方案问询

需求与示例 #

给定整数列表，需判断其中的非零数字是否构成单个连续组，示例如下：

[0, 0, -1, 2, -3, 0] → True
[1, 2, 3, 0, 0, 0] → True
[1, 2, 3, 4, 5, 0] → True
[1, 1, 1, 1, 1, 1] → True
[0, 1, 0, 3, 3, 3] → False
[1, 2, 3, 0, 3, 0] → False
[0, 0, 0, 0, 0, 0] → False

原有实现 #

你已写出如下实现代码：

def simple_contiguous_subset(data):
    if all(x == 0 for x in data): # 如果全零需要返回True可删除此行及下一行
        return False
    seen_zero_after_group = False
    in_non_zero_group = False
    for num in data:
        if num == 0:
            if not in_non_zero_group:
                seen_zero_after_group = False
            else:
                seen_zero_after_group = True
        else:
            if seen_zero_after_group:
                return False
            in_non_zero_group = True
    return True


data_list = [[0,0,-1,2,-3,0], [1,2,3,0,0,0], [1,2,3,4,5,0], [1,1,1,1,1,1], [0,1,0,3,3,3], [1,2,3,0,3,0], [0,0,0,0,0,0]]
for data in data_list:
    result = simple_contiguous_subset(data)
    print(f"{data} {result}")

更简洁高效的实现方式 #

方案1：利用itertools.groupby统计非零组数量

groupby可以按元素是否为零分组，核心逻辑是：非零组的数量必须恰好为1（全零则数量为0，返回False；多组则返回False）。

实现代码：

from itertools import groupby

def simple_contiguous_subset(data):
    # 统计所有非零分组的数量
    non_zero_group_count = sum(1 for is_zero, _ in groupby(data, key=lambda x: x == 0) if not is_zero)
    return non_zero_group_count == 1

方案2：优化版groupby实现（提前终止）

上面的代码会遍历所有分组，我们可以在计数超过1时直接返回，减少不必要的遍历：

from itertools import groupby

def simple_contiguous_subset(data):
    count = 0
    for is_zero, _ in groupby(data, key=lambda x: x == 0):
        if not is_zero:
            count += 1
            if count > 1:
                return False
    return count == 1

方案3：无需groupby的简洁遍历

通过两个状态变量记录状态，一次遍历完成判断，逻辑清晰且高效：

def simple_contiguous_subset(data):
    has_non_zero = False
    seen_zero_after_non_zero = False
    for num in data:
        if num != 0:
            if seen_zero_after_non_zero:
                return False
            has_non_zero = True
        else:
            if has_non_zero:
                seen_zero_after_non_zero = True
    return has_non_zero

这些实现都保留了原代码的O(n)时间复杂度，同时可读性更强、代码更简洁。其中方案2在遇到第二个非零组时直接返回，和原代码的提前终止逻辑一致，效率最优。

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