Design Underground System

mediumAmazon-tagged: Design-style coding33 stepsJava · Python · C++

Track check-ins and check-outs, and report average travel time between two stations. One map holds each traveller's open check-in (station, time). On check-out, add the trip to a route map keyed "from→to" holding total time and trip count.

▶ Watch the animated trace

Traced input: checkIn(45, Leyton, 3), checkIn(32, Paradise, 8), checkOut(45, Waterloo, 15), checkOut(32, Cambridge, 22), getAverageTime(Paradise, Cambridge), checkIn(10, Leyton, 24), checkOut(10, Waterloo, 38), getAverageTime(Leyton, Waterloo) → [null, null, null, null, null, 14.0, null, null, 13.0]. Original problem: leetcode.com/problems/design-underground-system/

Complexity

BestO(1)

AverageO(1)

WorstO(1)

SpaceO(P + S²)

P passengers in transit, S stations.

Key idea

Check-in data is only needed until that person checks out; remove it then.
Store totals and counts, not individual trips, so the average is O(1).
Routes are directional: A→B differs from B→A.

Common mistakes

Storing every trip and averaging on demand is O(trips) per query.
Integer division when averaging.

Solution

Java

class UndergroundSystem {
    private final Map<Integer, String> inStation = new HashMap<>();
    private final Map<Integer, Integer> inTime = new HashMap<>();
    private final Map<String, double[]> routes = new HashMap<>();

    void checkIn(int id, String station, int t) {
        inStation.put(id, station);
        inTime.put(id, t);
    }

    void checkOut(int id, String station, int t) {
        String key = inStation.remove(id) + "→" + station;
        int start = inTime.remove(id);
        double[] r = routes.computeIfAbsent(key, k -> new double[2]);
        r[0] += t - start;
        r[1]++;
    }

    double getAverageTime(String from, String to) {
        double[] r = routes.get(from + "→" + to);
        return r[0] / r[1];
    }
}

Python

class UndergroundSystem:
    def __init__(self):
        self.check_ins = {}
        self.routes = defaultdict(lambda: [0, 0])

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

    def checkOut(self, id: int, station: str, t: int) -> None:
        start, t0 = self.check_ins.pop(id)
        route = self.routes[(start, station)]
        route[0] += t - t0
        route[1] += 1

    def getAverageTime(self, startStation: str, endStation: str) -> float:
        total, count = self.routes[(startStation, endStation)]
        return total / count

C++

class UndergroundSystem {
    unordered_map<int, pair<string, int>> checkIns;
    map<pair<string, string>, pair<double, int>> routes;

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

    void checkOut(int id, string station, int t) {
        auto [start, t0] = checkIns[id];
        checkIns.erase(id);
        auto& r = routes[{start, station}];
        r.first += t - t0;
        r.second++;
    }

    double getAverageTime(string startStation, string endStation) {
        auto& r = routes[{startStation, endStation}];
        return r.first / r.second;
    }
};

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
Non-integer average (3 + 4) / 2.0.	`checkIn(1, A, 0), checkOut(1, B, 3), checkIn(2, A, 0), checkOut(2, B, 4), getAverageTime(A, B)`	`[null, null, null, null, null, 3.5]`
Same person, two trips The check-in record is cleared after each trip, and B→A is its own route.	`checkIn(7, A, 1), checkOut(7, B, 5), checkIn(7, B, 6), checkOut(7, A, 10), getAverageTime(B, A)`	`[null, null, null, null, null, 4.0]`