delay all remaining stations and check for nulls

2 files changed, 10 insertions, 2 deletions

diff --git a/src/Q10/Station.java b/src/Q10/Station.java
index a2b2ae2..0aa35b8 100644
--- a/src/Q10/Station.java
+++ b/src/Q10/Station.java
@@ -61,7 +61,8 @@ public class Station {
     this.setDepartureDate(advanceDate(this.getDepartureDate(), minutes));
 
     Calendar calendar = Calendar.getInstance();
-    calendar.setTime(this.getDepartureDate());
+    calendar.setTime(
+        this.getDepartureDate() != null ? this.getDepartureDate() : this.getArrivalDate());
     calendar.setTimeZone(TimeZone.getTimeZone("UTC"));
     this.setDay(calendar.get(Calendar.DAY_OF_MONTH));
   }
@@ -75,6 +76,8 @@ public class Station {
   }
 
   private Date advanceDate(Date original, int minutes) {
+    if (original == null) return null;
+
     return new Date(original.getTime() + (minutes * 60000));
   }
 }
diff --git a/src/Q10/TrainTimeTable.java b/src/Q10/TrainTimeTable.java
index 5883983..36e4294 100644
--- a/src/Q10/TrainTimeTable.java
+++ b/src/Q10/TrainTimeTable.java
@@ -21,7 +21,12 @@ public class TrainTimeTable {
     System.out.println();
 
     Station station = stationIn(city);
-    if (station != null) station.delayBy(minutes);
+    if (station != null) {
+      station.delayBy(minutes);
+      for (int i = schedule.indexOf(station) + 1; i < schedule.size(); i++) {
+        schedule.get(i).delayBy(minutes);
+      }
+    }
   }
 
   public void displaySchedule() {