Помогите оптимизировать код на Python

Задание:

Тимофей ищет место, чтобы построить себе дом. Улица, на которой он хочет жить, имеет длину n, то есть состоит из n одинаковых идущих подряд участков. Каждый участок либо пустой, либо на нём уже построен дом. Общительный Тимофей не хочет жить далеко от других людей на этой улице. Поэтому ему важно для каждого участка знать расстояние до ближайшего пустого участка. Если участок пустой, эта величина будет равна нулю — расстояние до самого себя. Помогите Тимофею посчитать искомые расстояния. Для этого у вас есть карта улицы. Дома в городе Тимофея нумеровались в том порядке, в котором строились, поэтому их номера на карте никак не упорядочены. Пустые участки обозначены нулями.

Формат ввода.

В первой строке дана длина улицы —– n (1 ≤ n ≤ 10^6). В следующей строке записаны n целых неотрицательных чисел — номера домов и обозначения пустых участков на карте (нули). Гарантируется, что в последовательности есть хотя бы один ноль. Номера домов (положительные числа) уникальны и не превосходят 10^9.

Формат вывода.

Для каждого из участков выведите расстояние до ближайшего нуля. Числа выводите в одну строку, разделяя их пробелами.

Пример 1

Ввод:
        
5
0 1 4 9 0

Вывод:

0 1 2 1 0

Пример 2

Ввод:
         
6            
0 7 9 4 8 20

Вывод:

0 1 2 3 4 5

Мой код:

from typing import List

def nearest_zero(arr: List[int]) -> List[int]:
arr.reverse()
zero_index: int = -1
for index, value in enumerate(arr):
    if value == 0:
        zero_index = index
    elif zero_index != -1 and (index - zero_index < value or value == -1):
        arr[index] = index - zero_index
return arr

def read_input() -> List[int]:
_ = input()
arr: list = [0 if x == '0' else -1 for x in input().split()]
return arr

def main() -> None:
array: list = read_input()
result: list = nearest_zero(nearest_zero(array))
print(*result)

if __name__ == '__main__':
    main()

То что нужно сделать: введите сюда описание изображения


Ответы (1 шт):

Автор решения: MBo

Как-то так:

def nearest_zero(a):
    zi = -1
    for i in range(len(a)):
        if a[i]:
            a[i] = i - zi
        else:
            for j in range(i - 1, -1 if zi < 0 else i - (i - zi) // 2 - 1, -1):             
                a[j] = i - j
            zi = i
    return a
→ Ссылка