Update Design Underground System.java

varunu28 · web-flow · commit 67036e7edd26 · 2022-03-22T00:58:58.000+05:30
diff --git a/Medium/Design Underground System.java b/Medium/Design Underground System.java
@@ -1,30 +1,33 @@
 class UndergroundSystem {
-  Map<Integer, String> checkInMap;
-  Map<String, Integer> timeMap;
-  Map<String, Integer> countMap;
+  
+  private Map<String, Integer> totalTimeFromStartToEndStation;
+  private Map<String, Integer> numberOfCustomersFromStartToEndStation;
+  private Map<Integer, Integer> customerIdToStartTime;
+  private Map<Integer, String> customerIdToStartStation;
+  
   public UndergroundSystem() {
-    checkInMap = new HashMap<>();
-    timeMap = new HashMap<>();
-    countMap = new HashMap<>();
+    this.totalTimeFromStartToEndStation = new HashMap<>();
+    this.numberOfCustomersFromStartToEndStation = new HashMap<>();
+    this.customerIdToStartTime = new HashMap<>();
+    this.customerIdToStartStation = new HashMap<>();
   }
 
   public void checkIn(int id, String stationName, int t) {
-    checkInMap.put(id, stationName + "," + t);
+    this.customerIdToStartTime.put(id, t);
+    this.customerIdToStartStation.put(id, stationName);
   }
 
   public void checkOut(int id, String stationName, int t) {
-    String val = checkInMap.get(id);
-    String prevStationName = val.split(",")[0];
-    int prevTime = Integer.parseInt(val.split(",")[1]);
-    String key = prevStationName + "|" + stationName;
-    timeMap.put(key, timeMap.getOrDefault(key, 0) + t - prevTime);
-    countMap.put(key, countMap.getOrDefault(key, 0) + 1);
+    int totalTime = t - this.customerIdToStartTime.get(id);
+    String key = this.customerIdToStartStation.get(id) + "|" + stationName;
+    this.totalTimeFromStartToEndStation.put(key, this.totalTimeFromStartToEndStation.getOrDefault(key, 0) + totalTime);
+    this.numberOfCustomersFromStartToEndStation.put(key, this.numberOfCustomersFromStartToEndStation.getOrDefault(key, 0) + 1);
   }
 
   public double getAverageTime(String startStation, String endStation) {
     String key = startStation + "|" + endStation;
-    return ((double) timeMap.get(key)) / countMap.get(key);
-  }
+    return this.totalTimeFromStartToEndStation.get(key) / (double) (this.numberOfCustomersFromStartToEndStation.get(key));
+   }
 }
 
 /**
