Car Fleet II LeetCode solution in Python

1
class Car:
2
    def __init__(self, pos, speed, idx, prev=None, next=None):
3
        self.pos = pos
4
        self.speed = speed
5
        self.idx = idx
6
        self.prev = prev
7
        self.next = next
8

9

10
class Solution:
11
    def getCollisionTimes(self, cars: List[List[int]]) -> List[float]:
12
        colis_times = [-1] * len(cars)
13
        cars = [Car(pos, sp, i) for i, (pos, sp) in enumerate(cars)]
14
        for i in range(len(cars) - 1):
15
            cars[i].next = cars[i + 1]
16
        for i in range(1, len(cars)):
17
            cars[i].prev = cars[i - 1]
18

19
        catchup_order = [
20
            ((b.pos - a.pos) / (a.speed - b.speed), a.idx, a)
21
            for i, (a, b) in enumerate(zip(cars, cars[1:]))
22
            if a.speed > b.speed
23
        ]
24
        heapify(catchup_order)
25

26
        while catchup_order:
27
            catchup_time, idx, car = heappop(catchup_order)
28
            if colis_times[idx] > -1:
29
                continue  # ith car has already caught up
30
            colis_times[idx] = catchup_time
31
            if not car.prev:
32
                continue  # no car is following us
33
            car.prev.next, car.next.prev = car.next, car.prev
34
            if car.next.speed >= car.prev.speed:
35
                continue  # the follower is too slow to catch up
36
            new_catchup_time = (car.next.pos - car.prev.pos) / (
37
                car.prev.speed - car.next.speed
38
            )
39
            heappush(catchup_order, (new_catchup_time, car.prev.idx, car.prev))
40

41
        return colis_times