leetcode/1512_design-underground-system/python3/solution.py

41 lines
1.3 KiB
Python
Raw Normal View History

2022-04-24 20:56:09 +00:00
# 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)