The machine must record two kinds of events and answer one question. A rider checks in at a station with a timestamp; later they check out at another station with another timestamp. Given a station pair, the machine must return the mean of all completed trips between them, including trips still in flight.
Two maps do the work. The first map holds active riders: rider id → (station, check-in time). When checkIn fires, the machine writes the entry, overwriting any prior entry for that id (a rider is never double-checked-in). When checkOut fires, the machine pops the active entry, computes elapsed = checkout time − check-in time, and folds it into the second map: (startStation, endStation) → [totalTime, tripCount]. The pair is keyed by a tuple or a flattened string like "Leyton,Waterloo" — string concatenation with a delimiter beats nested dicts for speed and clarity.
getAverageTime reads one entry from the second map and returns totalTime / tripCount. Everything is O(1) per operation; no scanning ever happens. The only arithmetic subtlety is floating point: the machine returns a float from integer division, so the average must be computed as totalTime / tripCount in Python 3 semantics, which yields a float automatically.
Edge cases: a rider checking out with no active check-in is a caller error the machine does not guard against; the spec guarantees valid interleavings. Multiple riders on the same route accumulate into the same bucket — order of check-out does not matter because the pair, not the rider, keys the bucket.
class UndergroundSystem:
def __init__(self):
self.active = {} # id -> (station, t)
self.trips = {} # (start, end) -> [total, count]
def checkIn(self, id, stationName, t):
self.active[id] = (stationName, t)
def checkOut(self, id, stationName, t):
start, t0 = self.active.pop(id)
key = (start, stationName)
total, count = self.trips.get(key, (0, 0))
self.trips[key] = (total + t - t0, count + 1)
def getAverageTime(self, startStation, endStation):
total, count = self.trips[(startStation, endStation)]
return total / count