Skip to content

1396. Design Underground System

  • Time: $O(1)$
  • Space: $O(|\texttt{passengers}| + |\texttt{stations}|^2)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
struct CheckIn {
  string stationName;
  int time;
};

struct CheckOut {
  int numTrips;
  int totalTime;
};

class UndergroundSystem {
 public:
  void checkIn(int id, string stationName, int t) {
    checkIns[id] = {stationName, t};
  }

  void checkOut(int id, string stationName, int t) {
    const auto [startStation, startTime] = checkIns[id];
    checkIns.erase(id);
    const string& route = startStation + "->" + stationName;
    ++checkOuts[route].numTrips;
    checkOuts[route].totalTime += t - startTime;
  }

  double getAverageTime(string startStation, string endStation) {
    const auto& [numTrips, totalTime] =
        checkOuts[startStation + "->" + endStation];
    return totalTime / (double)numTrips;
  }

 private:
  unordered_map<int, CheckIn> checkIns;       // {id: (stationName, time)}
  unordered_map<string, CheckOut> checkOuts;  // {route: (numTrips, totalTime)}
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class CheckIn {
  public String stationName;
  public int time;
  public CheckIn(String stationName, int time) {
    this.stationName = stationName;
    this.time = time;
  }
}

class CheckOut {
  public int numTrips;
  public int totalTime;
  public CheckOut(int numTrips, int totalTime) {
    this.numTrips = numTrips;
    this.totalTime = totalTime;
  }
}

class UndergroundSystem {
  public void checkIn(int id, String stationName, int t) {
    checkIns.put(id, new CheckIn(stationName, t));
  }

  public void checkOut(int id, String stationName, int t) {
    final CheckIn checkIn = checkIns.get(id);
    checkIns.remove(id);
    final String route = checkIn.stationName + "->" + stationName;
    checkOuts.putIfAbsent(route, new CheckOut(0, 0));
    ++checkOuts.get(route).numTrips;
    checkOuts.get(route).totalTime += t - checkIn.time;
  }

  public double getAverageTime(String startStation, String endStation) {
    final CheckOut checkOut = checkOuts.get(startStation + "->" + endStation);
    return checkOut.totalTime / (double) checkOut.numTrips;
  }

  private Map<Integer, CheckIn> checkIns = new HashMap<>();  // {id: (stationName, time)}
  private Map<String, CheckOut> checkOuts = new HashMap<>(); // {route: (numTrips, totalTime)}
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class UndergroundSystem:
  def __init__(self):
    # {id: (stationName, time)}
    self.checkIns = {}
    # {route: (numTrips, totalTime)}
    self.checkOuts = collections.defaultdict(lambda: [0, 0])

  def checkIn(self, id: int, stationName: str, t: int) -> None:
    self.checkIns[id] = (stationName, t)

  def checkOut(self, id: int, stationName: str, t: int) -> None:
    startStation, startTime = self.checkIns.pop(id)
    route = (startStation, stationName)
    self.checkOuts[route][0] += 1
    self.checkOuts[route][1] += t - startTime

  def getAverageTime(self, startStation: str, endStation: str) -> float:
    numTrips, totalTime = self.checkOuts[(startStation, endStation)]
    return totalTime / numTrips