Start learning free

LESSON 328 OF 569 · HASH TABLE

Design Underground System in Python: O(1) Averages with Hash Maps

Practice O(1) Read the full explanation

An underground railway wants to know how long its passengers travel between stations. Each passenger has a card with an id: they check in with it at the station where they start, and check out at the station where they leave. Design a system that records these trips and tells, for any two stations, the average time of the trips made so far from the first to the second.

Implement the class UndergroundSystem:

  • UndergroundSystem() creates an empty system.
  • checkIn(id, stationName, t): the passenger with the card id checks in at the station stationName at the time t. A passenger can be checked in at only one place at a time.
  • checkOut(id, stationName, t): the passenger with the card id checks out at the station stationName at the time t. This ends their trip: it went from the station where they checked in to stationName, and took t minus their check-in time.
  • getAverageTime(startStation, endStation) returns the average time of all the trips so far that went from startStation to endStation. Only finished trips count, and the direction matters: a trip from endStation to startStation belongs to a different route.

The input is a list of operations, applied in order to one new UndergroundSystem. The output lists what the getAverageTime calls return, in order, each with exactly two decimals. The starter code spells these names the way your language does.

 

Example 1:

Input:
    checkIn(45, "Leyton", 3)
    checkIn(32, "Paradise", 8)
    checkIn(27, "Leyton", 10)
    checkOut(45, "Waterloo", 15)
    checkOut(27, "Waterloo", 20)
    checkOut(32, "Cambridge", 22)
    getAverageTime("Paradise", "Cambridge")
    getAverageTime("Leyton", "Waterloo")
    checkIn(10, "Leyton", 24)
    getAverageTime("Leyton", "Waterloo")
Output: 14.00 11.00 11.00
Explanation:
Card 45 travels from Leyton to Waterloo in 15 - 3 = 12, card 27 on the same route in 20 - 10 = 10,
and card 32 from Paradise to Cambridge in 22 - 8 = 14.
The only trip from Paradise to Cambridge took 14. The two trips from Leyton to Waterloo average (12 + 10) / 2 = 11.
Card 10 has checked in at Leyton but not checked out yet: its trip doesn't count, and the average stays 11.

Example 2:

Input:
    checkIn(10, "Leyton", 3)
    checkOut(10, "Paradise", 8)
    getAverageTime("Leyton", "Paradise")
    checkIn(5, "Leyton", 10)
    checkOut(5, "Paradise", 16)
    getAverageTime("Leyton", "Paradise")
    checkIn(2, "Leyton", 21)
    checkOut(2, "Paradise", 30)
    checkIn(10, "Leyton", 32)
    checkOut(10, "Paradise", 37)
    getAverageTime("Leyton", "Paradise")
Output: 5.00 5.50 6.25
Explanation:
The trips from Leyton to Paradise take 5, 6, 9 and 5. After the first trip the average is 5, after two trips
(5 + 6) / 2 = 5.5, and after all four (5 + 6 + 9 + 5) / 4 = 6.25. Card 10 travels twice: each trip counts.

Example 3:

Input:
    checkIn(1, "A", 1)
    checkOut(1, "B", 4)
    checkIn(2, "B", 5)
    checkOut(2, "A", 12)
    checkIn(3, "A", 13)
    checkOut(3, "B", 15)
    getAverageTime("A", "B")
    getAverageTime("B", "A")
Output: 2.50 7.00
Explanation: The trips from A to B take 3 and 2, so their average is 2.5. The trip from B to A goes the other way: it belongs to its own route, with an average of 7.

 

Constraints:

  • 1 <= id, t <= 106
  • 1 <= stationName.length, startStation.length, endStation.length <= 10
  • Station names consist of uppercase and lowercase English letters and digits.
  • The calls come in time order: t never goes down from one call to the next, and every trip checks out later than it checked in, at a different station.
  • A passenger checks in only when they are not checked in already, and checks out only after checking in.
  • getAverageTime is only called for two stations with at least one finished trip between them.
  • All the trips together take at most 109, so every total fits in a 32-bit integer.
  • Every average in the tests has at most two decimals, so printing it with two decimals needs no rounding.
  • At most 105 calls in total are made to checkIn, checkOut and getAverageTime.

Note:

Each operation should run in O(1) time, apart from building the route's name from the two station names. Going through all the finished trips in getAverageTime would take O(n) time per call, where n is the number of trips so far.

STEP 1 OF 9
    Ctrl + Enter to check

    PYTHON · HASH TABLE

    Design Underground System in Python, explained

    Problem overview

    A railway wants the average travel time between two stations. A passenger holds a card with an id, checks in where the trip starts and checks out where it ends. The class UndergroundSystem has three methods:

    • checkIn(id, stationName, t) starts a trip for the card id. A card is checked in at only one place at a time.
    • checkOut(id, stationName, t) ends that trip. It went from the check-in station to stationName and took t minus the check-in time.
    • getAverageTime(startStation, endStation) returns the average time of all finished trips from the first station to the second.

    Only finished trips count, and direction matters: a trip from B to A belongs to a different route than one from A to B. The tests call getAverageTime only for routes that have at least one finished trip.

    Example: card 45 checks in at Leyton at time 3 and out at Waterloo at time 15, so the trip took 12. Card 27 makes the same trip from time 10 to time 20 and takes 10. The average for Leyton to Waterloo is (12 + 10) / 2 = 11.

    Approach: two hash maps and a running total

    The direct way is to record every finished trip as a record (start, end, duration) in one big list. getAverageTime then scans the whole list, picks the trips on the route and averages them. That costs O(n) per question, where n is the number of trips so far.

    An average is a total divided by a count, and both numbers can be kept up to date as trips finish. So the system stores, for each route, just two numbers: the total time and the number of trips. A check-out adds to them. A question is then one lookup and one division, and the individual trips never need to be stored.

    Two maps do the job:

    • self.checkIns: for each card id in use, the station and time of its check-in. A check-in only has to be remembered until the same card checks out, because only then is the trip's duration known.
    • self.routes: for each route, a pair (total time, number of trips).

    Every operation is a few dictionary lookups, so each is O(1) apart from building the route's name from the two station names.

    Step by step

    1. In __init__, set self.checkIns to an empty dictionary for the trips in progress.
    2. Set self.routes to an empty dictionary for the finished trips, per route. The route is written as one string, "start,end".
    3. In checkIn, store the tuple (stationName, t) under the key id. Nothing else is needed: the duration is unknown until check-out.
    4. In checkOut, take the passenger's check-in out with self.checkIns.pop(id) and unpack it into startStation and startTime. pop returns the stored value and removes the entry, which ends the trip in progress. The same card can check in again later, as card 10 does in the second example below.
    5. Build route from startStation, a comma and stationName. The order matters, since A to B and B to A are different routes. The comma keeps names apart: without it, "A" + "BC" and "AB" + "C" would both give "ABC".
    6. Unpack the route's pair into total and count with self.routes.get(route, (0, 0)). For a route with no finished trip yet, the default (0, 0) is used.
    7. Store the new pair under route: total + t - startTime and count + 1. A tuple can't be changed in place, so the new pair replaces the old one.
    8. In getAverageTime, unpack the pair stored for startStation + "," + endStation into total and count. The name is built exactly as in checkOut.
    9. Return total / count.

    The solution in Python

    class UndergroundSystem:
    def __init__(self):
        self.checkIns = {}
        self.routes = {}

    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
        total, count = self.routes.get(route, (0, 0))
        self.routes[route] = (total + t - startTime, count + 1)

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

    Here is the first example, with a look at the two maps along the way:

    u = UndergroundSystem()
u.checkIn(45, "Leyton", 3)
u.checkIn(32, "Paradise", 8)
u.checkIn(27, "Leyton", 10)
print(u.checkIns)
u.checkOut(45, "Waterloo", 15)
print(u.checkIns, u.routes)
u.checkOut(27, "Waterloo", 20)
u.checkOut(32, "Cambridge", 22)
print(u.routes)
print(u.getAverageTime("Paradise", "Cambridge"))
print(u.getAverageTime("Leyton", "Waterloo"))
u.checkIn(10, "Leyton", 24)
print(u.getAverageTime("Leyton", "Waterloo"))
print("%.2f" % u.getAverageTime("Leyton", "Waterloo"))

    {45: ('Leyton', 3), 32: ('Paradise', 8), 27: ('Leyton', 10)}
{32: ('Paradise', 8), 27: ('Leyton', 10)} {'Leyton,Waterloo': (12, 1)}
{'Leyton,Waterloo': (22, 2), 'Paradise,Cambridge': (14, 1)}
14.0
11.0
11.0
11.00

    After card 45 checks out, its entry leaves checkIns and the route holds one trip of 12. Card 10 checks in at the end, but its trip is unfinished and changes nothing. The tests print each average with two decimals, which the last line does with "%.2f".

    Python details worth knowing:

    • / always gives a float, so 7 / 2 is 3.5, while // would give 3 and lose the fraction. Averages need /.
    • a, b = pair unpacks a tuple into two names. It raises an error if the tuple has a different length.
    • dict.get(key, default) returns the default instead of raising a KeyError. In getAverageTime the code indexes with [...] on purpose, since the route is guaranteed to exist.
    • The parameter name id hides Python's built-in id() inside the method. The lesson uses it as given, and it does no harm here.

    Complexity

    Time: every method does a constant number of dictionary operations, so O(1) on average, plus the cost of building a route string, which is proportional to the length of the station names (at most 10 characters each). No operation depends on the number of trips.

    Space: O(P + R), where P is the number of passengers currently checked in and R is the number of distinct routes. Finished trips are not stored, only their totals.

    Edge cases and common mistakes

    • Direction. A to B and B to A are separate routes. After three trips the system holds {'A,B': (5, 2), 'B,A': (7, 1)} when two trips go from A to B (3 and 2 minutes) and one goes back (7 minutes).
    • An unfinished trip. It lives only in checkIns, so it never counts toward an average.
    • The same card twice. Because pop removes the check-in, the card can travel again, and each trip counts separately.
    • Integer division. total // count would give 5 instead of 5.5.
    • Forgetting the separator. Joining the station names without a comma can make two different routes collide.
    • Using [] for the route in checkOut. The first trip on a route would raise a KeyError. That is why checkOut uses get with a default.

    Testing

    Four trips on one route, printing the average after each:

    u = UndergroundSystem()
trips = [(10, 3, 8), (5, 10, 16), (2, 21, 30), (10, 32, 37)]
for card, start, end in trips:
    u.checkIn(card, "Leyton", start)
    u.checkOut(card, "Paradise", end)
    print("%.2f" % u.getAverageTime("Leyton", "Paradise"))

    5.00
5.50
6.67
6.25

    The durations are 5, 6, 9 and 5. Card 10 travels twice, and both trips count. The third average is 20 / 3, shown rounded by the format.

    Two stations in both directions:

    u = UndergroundSystem()
u.checkIn(1, "A", 1)
u.checkOut(1, "B", 4)
u.checkIn(2, "B", 5)
u.checkOut(2, "A", 12)
u.checkIn(3, "A", 13)
u.checkOut(3, "B", 15)
print(u.routes)
print(u.getAverageTime("A", "B"), u.getAverageTime("B", "A"))

    {'A,B': (5, 2), 'B,A': (7, 1)}
2.5 7.0
    Keep going with Python500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free