如何修改Python代码以支持任意数量0的列表最近0距离计算

用户提供的原代码仅能处理列表中只有一个0的情况，当存在多个0时，输出结果会出错。原代码如下：

lst = [2, 1, 0, 3, 5]
length = len(lst)
dist = []

for i in lst:
    if lst.index(0) > lst.index(i):
        diff = lst.index(0) - lst.index(i)
        dist.append(diff)
    elif lst.index(0) < lst.index(i):
        diff = abs(lst.index(i)) - abs(lst.index(0))
        dist.append(diff)
    elif lst.index(0) == lst.index(i):
        dist.append(0)
print(dist)

# Output: [2, 1, 0, 1, 2]

原代码的核心问题 #

lst.index(0)只会返回第一个出现的0的索引，当列表中有多个0时，后续的0以及它们附近的元素无法正确匹配到最近的0，导致距离计算错误。比如输入[0,2,0,4]，原代码会把第四个元素的距离错误计算为3，实际应该是1。

高效解决方案：两次遍历法 #

这种方法时间复杂度为O(n)，无需额外存储大量索引，效率最优：

def nearest_zero_distance(lst):
    length = len(lst)
    dist = [float('inf')] * length
    
    # 从左到右遍历，记录到左侧最近0的距离
    for i in range(length):
        if lst[i] == 0:
            dist[i] = 0
        elif i > 0:
            dist[i] = dist[i-1] + 1
    
    # 从右到左遍历，更新到右侧最近0的距离，取左右距离的最小值
    for i in range(length-2, -1, -1):
        dist[i] = min(dist[i], dist[i+1] + 1)
    
    return dist

# 测试示例
test_lst1 = [2, 1, 0, 3, 5]
print(nearest_zero_distance(test_lst1))  # 输出: [2, 1, 0, 1, 2]

test_lst2 = [0, 2, 0, 4, 1, 0]
print(nearest_zero_distance(test_lst2))  # 输出: [0, 1, 0, 1, 1, 0]

test_lst3 = [5, 4, 3, 2, 1, 0]
print(nearest_zero_distance(test_lst3))  # 输出: [5, 4, 3, 2, 1, 0]

直观解决方案：收集0的索引再计算 #

如果追求逻辑简单，可先收集所有0的位置，再逐个计算每个元素到这些0的最小距离：

def nearest_zero_distance(lst):
    zero_indices = [i for i, val in enumerate(lst) if val == 0]
    dist = []
    for idx in range(len(lst)):
        # 计算当前索引到每个0的距离，取最小值
        min_dist = min(abs(idx - z_idx) for z_idx in zero_indices)
        dist.append(min_dist)
    return dist

# 测试示例
test_lst = [1, 0, 3, 0, 5]
print(nearest_zero_distance(test_lst))  # 输出: [1, 0, 1, 0, 1]

这种方法逻辑易懂，但时间复杂度为O(n*m)（m为0的数量），当列表极大且0较多时，效率不如两次遍历法。

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