Design Underground System (LeetCode 1396)

The problem

Design a system that tracks passengers on a metro network:

• checkIn(id, station, t): passenger id enters at station at time t;
• checkOut(id, station, t): the same passenger leaves at station at time t;
• getAverageTime(start, end): return the average travel time over all completed journeys that went directly from start to end.

Each passenger is checked in to at most one station at a time, and every query asks about a route with at least one completed trip.

checkIn(45, "A", 3);  checkOut(45, "B", 8)     trip A->B took 5
checkIn(27, "A", 5);  checkOut(27, "B", 12)    trip A->B took 7
getAverageTime("A", "B") -> 6.0                (5 + 7) / 2

How to think about it

Start with brute force. Log every completed trip as (start, end, duration) in one list. getAverageTime scans the whole list, keeps the matching trips and averages them: O(total trips) per query. With many queries this is slow, and you store data you never need again.

The key insight: an average only needs a sum and a count. You don’t need to remember individual trips at all. For each route (start, end), keep a running total of travel time and a count of trips. Each checkout adds one trip to its route in O(1), and the average is just total / count, also O(1).

The second piece is the open journeys. At checkout you need to know where and when this passenger entered. So keep a second map, id → (station, time), for passengers currently inside the network. On checkout, pop their entry (they are no longer travelling), compute the duration, and fold it into the route totals.

Which pattern is this? Design a data structure with two hash maps, one for in-progress state and one for aggregates. The tell is “answer averages or totals over a stream of events”: maintain the aggregate incrementally instead of recomputing it.

Approach

1. Keep checkins: id → (station, t), and totals: (start, end) → [total_time, count].
2. checkIn: store (station, t) under id.
3. checkOut: pop id from checkins to get the start station and time. Add t - start_time to totals[(start, station)][0] and 1 to its count.
4. getAverageTime: return total / count for the route.

Solution

from collections import defaultdict

class UndergroundSystem:
    def __init__(self):
        self.checkins = {}; self.totals = defaultdict(lambda: [0, 0])   # route -> [total, count]
    def checkIn(self, uid, station, t): self.checkins[uid] = (station, t)
    def checkOut(self, uid, station, t):
        s, st = self.checkins.pop(uid); key = (s, station)              # journey is now closed
        self.totals[key][0] += t - st; self.totals[key][1] += 1
    def getAverageTime(self, start, end):
        tot, cnt = self.totals[(start, end)]; return tot / cnt
# checkIn(45,A,3) checkOut(45,B,8) checkIn(27,A,5) checkOut(27,B,12) getAverageTime(A,B)->6.0

Dry run

The example, plus a trip in the opposite direction to show routes are directional:

A→B and B→A are separate keys, so their averages never mix.

Complexity

• Time: O(1) average per operation. Each is a constant number of hash-map reads and writes.
• Space: O(P + S²) in the worst case: P passengers currently travelling, plus one entry per distinct (start, end) pair that has been travelled.

Edge cases and pitfalls

• Routes are directional. Use the ordered pair (start, end) as the key, not a set of the two stations.
• Pop the check-in on checkout. Leaving it in the map leaks memory, and a passenger can check in again later with the same id.
• Integer division. In Python 3, / is true division; in Java, cast to double before dividing.
• Unknown route would divide by zero. The problem promises at least one trip; mention a guard if asked.
• Storing every trip works but makes queries O(trips). Keep only the aggregate.

Related problems

• Design Twitter (LeetCode 355): another design built from a few hash maps keyed by user.
• Time Based Key-Value Store (LeetCode 981): a map whose values carry timestamps, queried later.
• Design HashMap (LeetCode 706): what makes all those O(1) map operations possible.