41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
# Time: O(N)
|
|
# Space: O(P) P is num of passengers who checkin at same time worst case
|
|
# + O(S^2) S is num of stations, pair up every station so S^2 permutation
|
|
|
|
class UndergroundSystem:
|
|
|
|
def __init__(self):
|
|
self.checkins = {}
|
|
self.totals = {}
|
|
|
|
def checkIn(self, id: int, stationName: str, t: int) -> None:
|
|
self.checkins[id] = (stationName, t)
|
|
|
|
def checkOut(self, id: int, endStation: str, end_t: int) -> None:
|
|
startStation, start_t = self.checkins[id]
|
|
|
|
key = (startStation, endStation)
|
|
duration = end_t - start_t
|
|
|
|
curr_total = self.totals.get(key, (0, 0))
|
|
|
|
# Increment total and count of journeys
|
|
self.totals[key] = (curr_total[0] + duration, curr_total[1] + 1)
|
|
|
|
# Journey is over for this passenger, we can remove it since we
|
|
# tallied the data
|
|
del self.checkins[id]
|
|
|
|
|
|
def getAverageTime(self, startStation: str, endStation: str) -> float:
|
|
total_duration, num_trips = self.totals[(startStation, endStation)]
|
|
|
|
return total_duration / num_trips
|
|
|
|
|
|
# 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)
|