1396. Design Underground System

Problem

Implement the UndergroundSystem class:

void checkIn(int id, string stationName, int t)
- A customer with a card id equal to id, gets in the station stationName at time t.
- A customer can only be checked into one place at a time.
void checkOut(int id, string stationName, int t)
- A customer with a card id equal to id, gets out from the station stationName at time t.
double getAverageTime(string startStation, string endStation)
- Returns the average time to travel between the startStation and the endStation.
- The average time is computed from all the previous traveling from startStation to endStation that happened directly.
- Call to getAverageTime is always valid.

You can assume all calls to checkIn and checkOut methods are consistent. If a customer gets in at time t1 at some station, they get out at time t2 with t2 > t1. All events happen in chronological order.

Example 1:

Input
["UndergroundSystem","checkIn","checkIn","checkIn","checkOut","checkOut","checkOut","getAverageTime","getAverageTime","checkIn","getAverageTime","checkOut","getAverageTime"]
[[],[45,"Leyton",3],[32,"Paradise",8],[27,"Leyton",10],[45,"Waterloo",15],[27,"Waterloo",20],[32,"Cambridge",22],["Paradise","Cambridge"],["Leyton","Waterloo"],[10,"Leyton",24],["Leyton","Waterloo"],[10,"Waterloo",38],["Leyton","Waterloo"]]

Output
[null,null,null,null,null,null,null,14.00000,11.00000,null,11.00000,null,12.00000]

Explanation
UndergroundSystem undergroundSystem = new UndergroundSystem();
undergroundSystem.checkIn(45, "Leyton", 3);
undergroundSystem.checkIn(32, "Paradise", 8);
undergroundSystem.checkIn(27, "Leyton", 10);
undergroundSystem.checkOut(45, "Waterloo", 15);
undergroundSystem.checkOut(27, "Waterloo", 20);
undergroundSystem.checkOut(32, "Cambridge", 22);
undergroundSystem.getAverageTime("Paradise", "Cambridge");       // return 14.00000. There was only one travel from "Paradise" (at time 8) to "Cambridge" (at time 22)
undergroundSystem.getAverageTime("Leyton", "Waterloo");          // return 11.00000. There were two travels from "Leyton" to "Waterloo", a customer with id=45 from time=3 to time=15 and a customer with id=27 from time=10 to time=20. So the average time is ( (15-3) + (20-10) ) / 2 = 11.00000
undergroundSystem.checkIn(10, "Leyton", 24);
undergroundSystem.getAverageTime("Leyton", "Waterloo");          // return 11.00000
undergroundSystem.checkOut(10, "Waterloo", 38);
undergroundSystem.getAverageTime("Leyton", "Waterloo");          // return 12.00000

Example 2:

Input
["UndergroundSystem","checkIn","checkOut","getAverageTime","checkIn","checkOut","getAverageTime","checkIn","checkOut","getAverageTime"]
[[],[10,"Leyton",3],[10,"Paradise",8],["Leyton","Paradise"],[5,"Leyton",10],[5,"Paradise",16],["Leyton","Paradise"],[2,"Leyton",21],[2,"Paradise",30],["Leyton","Paradise"]]

Output
[null,null,null,5.00000,null,null,5.50000,null,null,6.66667]

Explanation
UndergroundSystem undergroundSystem = new UndergroundSystem();
undergroundSystem.checkIn(10, "Leyton", 3);
undergroundSystem.checkOut(10, "Paradise", 8);
undergroundSystem.getAverageTime("Leyton", "Paradise"); // return 5.00000
undergroundSystem.checkIn(5, "Leyton", 10);
undergroundSystem.checkOut(5, "Paradise", 16);
undergroundSystem.getAverageTime("Leyton", "Paradise"); // return 5.50000
undergroundSystem.checkIn(2, "Leyton", 21);
undergroundSystem.checkOut(2, "Paradise", 30);
undergroundSystem.getAverageTime("Leyton", "Paradise"); // return 6.66667

Constraints:

There will be at most 20000 operations.
1 <= id, t <= 106
All strings consist of uppercase and lowercase English letters, and digits.
1 <= stationName.length <= 10
Answers within 105 of the actual value will be accepted as correct.

Problem Summary

Please implement a class UndergroundSystem that supports the following 3 methods:

1. checkIn(int id, string stationName, int t)
- A passenger with card id id enters the subway station stationName at time t.
- A passenger can only enter or leave one subway station at the same time.
2. checkOut(int id, string stationName, int t)
- A passenger with card id id leaves the subway station stationName at time t.
3. getAverageTime(string startStation, string endStation)
- Return the average time taken to travel from subway station startStation to subway station endStation.
- The trips used to calculate the average time include all trips so far that went directly from startStation to endStation.
- When getAverageTime is called, the queried route contains at least one trip.

You can assume that all calls to checkIn and checkOut are logically valid. That is, if a customer arrives at some subway station at time t1, then their departure time t2 must satisfy t2 > t1. All events are given in chronological order.

Solution Approach

Maintain 2 maps. One mapA stores the mapping between passenger id and (check-in time, station name). The other mapB stores the relationship between the start station and end station with the total time spent and total number of people. Whenever someone checkin(), update the information in mapA. Whenever someone checkout(), update the information in mapB, and delete the key-value pair corresponding to the passenger id in mapA. Finally, when calling the getAverageTime() function, simply calculate based on the information stored in mapB.

Code

package leetcode

type checkin struct {
	station string
	time    int
}

type stationTime struct {
	sum, count float64
}

type UndergroundSystem struct {
	checkins     map[int]*checkin
	stationTimes map[string]map[string]*stationTime
}

func Constructor() UndergroundSystem {
	return UndergroundSystem{
		make(map[int]*checkin),
		make(map[string]map[string]*stationTime),
	}
}

func (s *UndergroundSystem) CheckIn(id int, stationName string, t int) {
	s.checkins[id] = &checkin{stationName, t}
}

func (s *UndergroundSystem) CheckOut(id int, stationName string, t int) {
	checkin := s.checkins[id]
	destination := s.stationTimes[checkin.station]
	if destination == nil {
		s.stationTimes[checkin.station] = make(map[string]*stationTime)
	}
	st := s.stationTimes[checkin.station][stationName]
	if st == nil {
		st = new(stationTime)
		s.stationTimes[checkin.station][stationName] = st
	}
	st.sum += float64(t - checkin.time)
	st.count++
	delete(s.checkins, id)
}

func (s *UndergroundSystem) GetAverageTime(startStation string, endStation string) float64 {
	st := s.stationTimes[startStation][endStation]
	return st.sum / st.count
}

/**
 * Your UndergroundSystem object will be instantiated and called as such:
 * obj := Constructor();
 * obj.CheckIn(id,stationName,t);
 * obj.CheckOut(id,stationName,t);
 * param_3 := obj.GetAverageTime(startStation,endStation);
 */