class UndergroundSystem:
def __init__(self):
self.station_id_dict = {}
# 两个int 字典,一个记录总时长,一个记录总个数,可以进一步优化空间
self.station_duration_dict = collections.defaultdict(list)
def checkIn(self, id: int, stationName: str, t: int) -> None:
self.station_id_dict[id] = (stationName, t)
def checkOut(self, id: int, stationName: str, t: int) -> None:
in_station, in_t = self.station_id_dict.pop(id)
self.station_duration_dict[(in_station, stationName)].append(t - in_t)
def getAverageTime(self, startStation: str, endStation: str) -> float:
time_list = self.station_duration_dict[(startStation, endStation)]
return sum(time_list) / len(time_list)