# 1396. Design Underground System

{% embed url="<https://leetcode.com/problems/design-underground-system/>" %}

乏善可陈的design题，利用`hashmap`解决，问题主要是动态计算均值。

### 1/3/2021 Update:

第二次写，思路还是一样，多用了一个dict，数据结构清楚一点，而且学会了pop掉下车的乘客，节约空间。

{% tabs %}
{% tab title="Python" %}

```python
class UndergroundSystem:

    def __init__(self):
        self.trips = collections.defaultdict(lambda : (0,0))
        self.passengers = dict()
        

    def checkIn(self, id: int, stationName: str, t: int) -> None:
        self.passengers[id] = (stationName, t)

    def checkOut(self, id: int, stationName: str, t: int) -> None:
        trip = (self.passengers[id][0], stationName)
        newCost = t - self.passengers[id][1]
        formerCost, formerCount = self.trips[trip][0], self.trips[trip][1]
        self.trips[trip] = ((formerCost*formerCount+newCost)/(formerCount+1), formerCount+1)

    def getAverageTime(self, startStation: str, endStation: str) -> float:
        return self.trips[(startStation, endStation)][0]


# Your UndergroundSystem object will be instantiated and called as such:
# obj = UndergroundSystem()
# obj.checkIn(id,stationName,t)
# obj.checkOut(id,stationName,t)
# param_3 = obj.getAverageTime(startStation,endStation)
```

{% endtab %}

{% tab title="3 dict version" %}

```python
class UndergroundSystem:

    def __init__(self):
        self.numCustomer = collections.defaultdict(int)
        self.aveTime = collections.defaultdict(int)
        self.curr = dict()

    def checkIn(self, id: int, stationName: str, t: int) -> None:
        self.curr[id] = (stationName, t)

    def checkOut(self, id: int, stationName: str, t: int) -> None:
        route = (self.curr[id][0], stationName)
        newTime = t - self.curr[id][1]
        formerTime, formerCustomer = self.aveTime[route], self.numCustomer[route]
        self.aveTime[route] = ((formerTime * formerCustomer) + newTime) / (formerCustomer + 1)
        self.numCustomer[route] += 1
        self.curr.pop(id)

    def getAverageTime(self, startStation: str, endStation: str) -> float:
        return self.aveTime[(startStation, endStation)]


# Your UndergroundSystem object will be instantiated and called as such:
# obj = UndergroundSystem()
# obj.checkIn(id,stationName,t)
# obj.checkOut(id,stationName,t)
# param_3 = obj.getAverageTime(startStation,endStation)
```

{% endtab %}
{% endtabs %}
