Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import heapq
- class UndergroundSystem:
- def __init__(self):
- self._start = {} # maps trip start to customer id and time
- self._end = {} # maps trip end to customer id and time
- def checkIn(self, id: int, stationName: str, t: int) -> None:
- start = self._start.get(stationName, [])
- heapq.heappush(start, (id, t))
- self._start[stationName] = start
- def checkOut(self, id: int, stationName: str, t: int) -> None:
- end = self._end.get(stationName, [])
- heapq.heappush(end, (id, t))
- self._end[stationName] = end
- def getAverageTime(self, startStation: str, endStation: str) -> float:
- start_pq = self._start[startStation].copy()
- end_pq = self._end[endStation].copy()
- trips = []
- while start_pq:
- customer, start_time = heapq.heappop(start_pq)
- while end_pq and end_pq[0][0] < customer:
- heapq.heappop(end_pq)
- if end_pq and end_pq[0][0] == customer:
- trips.append(end_pq[0][1] - start_time)
- heapq.heappop(end_pq)
- return sum(trips) / len(trips)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement