1396. Design Underground System - LeetCode

from collections import defaultdict


class UndergroundSystem:
    def __init__(self):
        checkInStationNameT = checkOutStationNameT = str
        checkInTimeT = travelTimeT = int
        idT = countT = int
        self.check_in: dict[idT, tuple[checkInTimeT, checkInStationNameT]] = {}
        self.time: defaultdict[
            tuple[checkInStationNameT, checkOutStationNameT], tuple[travelTimeT, countT]
        ] = defaultdict(lambda: (0, 0))

    def checkIn(self, id: int, stationName: str, t: int) -> None:
        self.check_in[id] = (t, stationName)

    def checkOut(self, id: int, stationName: str, t: int) -> None:
        check_in_time, check_in_station_name = self.check_in.pop(id)
        _key = check_in_station_name, stationName
        _travel_time, _count = self.time[_key]
        self.time[_key] = _travel_time + (t - check_in_time), _count + 1

    def getAverageTime(self, startStation: str, endStation: str) -> float:
        _key = startStation, endStation
        travel_time, count = self.time[_key]
        return travel_time / count