Design Underground System

Problem

Track travel times between stations on a subway network. Implement void checkIn(int id, String stationName, int t) — customer id enters at stationName at time t; void checkOut(int id, String stationName, int t) — customer id exits at stationName at time t; double getAverageTime(String startStation, String endStation) — the average travel time of all customers who went directly from startStation to endStation.

Constraints

• There is always a matching checkIn before each checkOut for an id.
• getAverageTime is only queried for a route that has at least one completed trip.
• at most 2·10⁴ calls

Approach — Hashing & Counting

This is a Hashing & Counting problem. The idea: store what you've seen in a hash map for O(1) lookups, trading a little space to avoid a nested scan. Work through the reference code below line by line, then re-derive it yourself in the editor — that's how the pattern sticks.

Complexity: O(n) time, O(n) space.

Solution code

Python

class UndergroundSystem:
    def __init__(self):
        self.checked = {}      # id -> (station, time)
        self.totals = {}       # (start, end) -> [total_time, count]

    def checkIn(self, id, stationName, t):
        self.checked[id] = (stationName, t)

    def checkOut(self, id, stationName, t):
        start, t0 = self.checked.pop(id)
        key = (start, stationName)
        rec = self.totals.setdefault(key, [0, 0])
        rec[0] += t - t0
        rec[1] += 1

    def getAverageTime(self, startStation, endStation):
        total, count = self.totals[(startStation, endStation)]
        return total / count

Java

import java.util.HashMap;
import java.util.Map;

class UndergroundSystem {
    // id -> [startStation, checkInTime]
    private Map<Integer, Object[]> checkedIn;
    // "start->end" -> [totalTime, tripCount]
    private Map<String, double[]> routeStats;

    public UndergroundSystem() {
        checkedIn = new HashMap<>();
        routeStats = new HashMap<>();
    }

    public void checkIn(int id, String stationName, int t) {
        checkedIn.put(id, new Object[]{stationName, t});
    }

    public void checkOut(int id, String stationName, int t) {
        Object[] entry = checkedIn.remove(id);
        String start = (String) entry[0];
        int startTime = (int) entry[1];
        String route = start + "->" + stationName;
        double[] stats = routeStats.computeIfAbsent(route, k -> new double[2]);
        stats[0] += t - startTime;
        stats[1] += 1;
    }

    public double getAverageTime(String startStation, String endStation) {
        double[] stats = routeStats.get(startStation + "->" + endStation);
        return stats[0] / stats[1];
    }
}

Practice it

Reading a solution isn't the same as being able to write it under pressure. Open this problem in the in-browser editor, hide the solution, and solve it from scratch — your code runs against real test cases instantly.