1396. Design Underground System

class UndergroundSystem:

    def __init__(self):
        self.data = {}
        self.pending = {}

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

    def checkOut(self, id: int, stationName: str, t: int) -> None:
        prevStation = self.pending[id][0]
        prevTime = self.pending[id][1]
        diff = t - prevTime
        if (prevStation, stationName) in self.data:
            data = self.data[(prevStation, stationName)]
            newCount = data[0] + 1
            total = data[1] + diff
            self.data[(prevStation, stationName)] = [newCount, total]
        else:
            self.data[(prevStation, stationName)] = [1, diff]

    def getAverageTime(self, startStation: str, endStation: str) -> float:
        return self.data[(startStation, endStation)][1] / self.data[(startStation, endStation)][0]