Design Underground System

https://leetcode.com/problems/design-underground-system/description/

class UndergroundSystem:

    def __init__(self):
        self.checkInMap = {} # id -> (startStation, time)
        self.totalMap = {} # (start, end) -> [totalTime, count]

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

    def checkOut(self, id: int, stationName: str, t: int) -> None:
        start, time = self.checkInMap[id]
        route = (start, stationName)
        if route not in self.totalMap:
            self.totalMap[route] = [0,0]
        self.totalMap[route][0] += t - time
        self.totalMap[route][1] += 1


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

key is to use hashmap and when check out we calculate overall time that customer have traveled.

at the end just we can just return total / count 

overall complexity is very efficient due to hash map