Advertisement
Guest User

Untitled

a guest
Mar 28th, 2020
91
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.20 KB | None | 0 0
  1. import heapq
  2.  
  3.  
  4. class UndergroundSystem:
  5.  
  6.     def __init__(self):
  7.         self._start = {}  # maps trip start to customer id and time
  8.         self._end = {}  # maps trip end to customer id and time
  9.        
  10.  
  11.     def checkIn(self, id: int, stationName: str, t: int) -> None:
  12.         start = self._start.get(stationName, [])
  13.         heapq.heappush(start, (id, t))
  14.         self._start[stationName] = start
  15.        
  16.  
  17.     def checkOut(self, id: int, stationName: str, t: int) -> None:
  18.         end = self._end.get(stationName, [])
  19.         heapq.heappush(end, (id, t))
  20.         self._end[stationName] = end
  21.        
  22.  
  23.     def getAverageTime(self, startStation: str, endStation: str) -> float:
  24.         start_pq = self._start[startStation].copy()
  25.         end_pq = self._end[endStation].copy()
  26.        
  27.         trips = []
  28.         while start_pq:
  29.             customer, start_time = heapq.heappop(start_pq)
  30.             while end_pq and end_pq[0][0] < customer:
  31.                 heapq.heappop(end_pq)
  32.            
  33.             if end_pq and end_pq[0][0] == customer:
  34.                 trips.append(end_pq[0][1] - start_time)
  35.                 heapq.heappop(end_pq)    
  36.        
  37.         return sum(trips) / len(trips)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement