Помогите оптимизировать код на 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 шт):
Как-то так:
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
