package leetcode;

import java.util.HashMap;

class UndergroundSystem {
  HashMap<Integer, Checkin> checkins = new HashMap<>();
  HashMap<String, HashMap<String, Stats>> statsLookup = new HashMap<>();

  public void checkIn(int id, String stationName, int t) {
    checkins.put(id, new Checkin(stationName, t));
  }

  public void checkOut(int id, String stationName, int t) {
    Checkin checkin = checkins.get(id);
    statsLookup.put(checkin.stationIn, statsLookup.getOrDefault(checkin.stationIn, new HashMap<>()));
    HashMap<String, Stats> inner = statsLookup.get(checkin.stationIn);
    inner.put(stationName, inner.getOrDefault(stationName, new Stats()));
    inner.get(stationName).samples++;
    inner.get(stationName).totalTime += t - checkin.timeIn;
  }

  public double getAverageTime(String startStation, String endStation) {
    return statsLookup.get(startStation).get(endStation).getAverage();
  }

  static class Stats {
    int totalTime = 0;
    int samples = 0;

    public double getAverage() {
      return totalTime / (samples + 0.0);
    }
  }

  static class Checkin {
    String stationIn;
    int timeIn;

    public Checkin(String s, int t) {
      this.stationIn = s;
      this.timeIn = t;
    }
  }
}