underground-system py3

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

@@ -0,0 +1,40 @@
+# 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)