Design Underground System

class UndergroundSystem:
 
    def __init__(self):
        self.user_checked_in_time = defaultdict(float)
        self.user_checked_in_station = defaultdict(str)
        self.start_end_counter = defaultdict(float)
        self.start_end_sum = defaultdict(float)
        self.checked_in_user_set = set()
        
 
    def checkIn(self, id: int, stationName: str, t: int) -> None:
        if id not in self.checked_in_user_set:
            self.checked_in_user_set.add(id)
            self.user_checked_in_time[id] = t
            self.user_checked_in_station[id] = stationName
        else:
            print("User already checked in!")
 
    def checkOut(self, id: int, stationName: str, t: int) -> None:
        if id not in self.checked_in_user_set:
            print("User not checked in previously")
            return
        
        start_end_tuple = (self.user_checked_in_station[id], stationName)
        start_end_time = float(t) - self.user_checked_in_time[id]
        self.start_end_counter[start_end_tuple] += 1
        self.start_end_sum[start_end_tuple] += start_end_time
        
        # Cleanup #O(N)
        self.checked_in_user_set.remove(id)
        self.user_checked_in_station.pop(id)
        self.user_checked_in_time.pop(id)
        
    def getAverageTime(self, startStation: str, endStation: str) -> float:
        start_end_tuple = (startStation, endStation)
        return self.start_end_sum[start_end_tuple] / self.start_end_counter[start_end_tuple]
        
# Your UndergroundSystem object will be instantiated and called as such:
# obj = UndergroundSystem()
# obj.checkIn(id,stationName,t)
# obj.checkOut(id,stationName,t)
# param_3 = obj.getAverageTime(startStation,endStation)

---

References

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