This is a very famous Bloomberg data structure design question (I got this question in my phone screen!).
Solution
The solution is rather straightforward. We need one hashmap to store the current check ins which have not yet checked out. Then, we will use another hashmap of hashmaps like this: unordered_map<string, unordered_map<string, pair<int,int>>>. The key of this hashmap is the starting station and for each starting station the map stores ending stations which then have the information stored in a pair. Pair’s first element is the total time and second element is total number of people who travelled this path: unordered_map<starting_station, unordered_map<ending_station, pair<total_time_taken, number_of_travellers>>>. This enables us to have O(1) time complexity for all operations.
class UndergroundSystem {
public:
unordered_map<string, unordered_map<string, pair<int,int>>> map;
unordered_map<int, pair<string,int>> checkIns;
UndergroundSystem() {
}
void checkIn(int id, string stationName, int t) {
checkIns[id] = {stationName, t};
}
void checkOut(int id, string stationName, int t) {
auto checkIn = checkIns[id];
checkIns.erase(id);
map[checkIn.first][stationName].first += t-checkIn.second;
map[checkIn.first][stationName].second += 1;
}
double getAverageTime(string startStation, string endStation) {
return (double) map[startStation][endStation].first/(double)map[startStation][endStation].second;
}
};
Some space can further be conserved by a little trick. Instead of having a 2D map of starting_station -> map(ending_stations) we can have a map which has the key “startStation_endStation”. Let’s call this map checkoutMap.
class UndergroundSystem {
public:
unordered_map<string, pair<int, int>> checkoutMap; // Route - {TotalTime, Count}
unordered_map<int, pair<string, int>> checkInMap; // Uid - {StationName, Time}
UndergroundSystem() {}
void checkIn(int id, string stationName, int t) {
checkInMap[id] = {stationName, t};
}
void checkOut(int id, string stationName, int t) {
auto& checkIn = checkInMap[id];
string route = checkIn.first + "_" + stationName;
checkoutMap[route].first += t - checkIn.second;
checkoutMap[route].second += 1;
}
double getAverageTime(string startStation, string endStation) {
string route = startStation + "_" + endStation;
auto& checkout = checkoutMap[route];
return (double) checkout.first / checkout.second;
}
};