Problem statement

https://binarysearch.com/problems/Underground-Tunnel/

Solution

Equal to Leetcode 1396. Design Underground System.

Complexity

Time compexlty is O(1) for all 3 operations. Space complexity potentially is O(Q), where Q is total number of queries.

Code

class UndergroundTunnel:
    def __init__(self):
        self.ids = {}
        self.pairs = Counter()
        self.freqs = Counter()
        
    def checkIn(self, id, stationName, t):
        self.ids[id] = (stationName, t)

    def checkOut(self, id, stationName, t):
        Name2, t2 = self.ids.pop(id)
        self.pairs[(Name2, stationName)] += t-t2
        self.freqs[(Name2, stationName)] += 1
        
    def averageTime(self, startStation, endStation):
        return self.pairs[startStation, endStation]/self.freqs[startStation, endStation]