Design Underground System LeetCode solution in Python

1
# Runtime: 441 ms (Top 30.19%) | Memory: 23.9 MB (Top 87.19%)
2
from collections import defaultdict
3

4

5
class UndergroundSystem:
6
    def __init__(self):
7
        self.checkin = {}
8
        self.traveltime = defaultdict(dict)
9

10
    def checkIn(self, id: int, stationName: str, t: int) -> None:
11
        self.checkin[id] = (stationName, t)
12

13
    def checkOut(self, id: int, stationName: str, t: int) -> None:
14
        startStation, startTime = self.checkin[id]
15
        del self.checkin[id]
16
        if stationName not in self.traveltime[startStation]:
17
            self.traveltime[startStation][stationName] = []
18
        self.traveltime[startStation][stationName].append(t - startTime)
19

20
    def getAverageTime(self, startStation: str, endStation: str) -> float:
21
        return sum(self.traveltime[startStation][endStation]) / len(
22
            self.traveltime[startStation][endStation]
23
        )