diff --git a/src/ShortestPath/Transit/TransitShortestPathContainer.cpp b/src/ShortestPath/Transit/TransitShortestPathContainer.cpp
index 0c35c86dc57b9a93d1ae2b9b8f91e3b50130ff0b..bee8ce44b11ed7015d414e6a774d9e77802dac4f 100644
--- a/src/ShortestPath/Transit/TransitShortestPathContainer.cpp
+++ b/src/ShortestPath/Transit/TransitShortestPathContainer.cpp
@@ -3,6 +3,7 @@
 //
 
 #include "TransitShortestPathContainer.h"
+#include "TransitStateContainer.h"
 
 void TransitShortestPathContainer::addShortestPathCollection(int startNodeIndex,
                                                              const std::pair<int, std::vector<TransitShortestPath>>& shortestPathList) {
@@ -10,13 +11,13 @@ void TransitShortestPathContainer::addShortestPathCollection(int startNodeIndex,
 }
 
 void TransitShortestPathContainer::addShortestPathCollection(int startNodeIndex, int startingInstant, int graphSize,
-                                                             const std::vector<TransitAlgorithmState>& algorithmResultStates) {
+                                                             const TransitStateContainer& algorithmResultStates) {
     std::vector<TransitShortestPath> shortestPathList;
     shortestPathList.reserve(graphSize);
 
     //Convert states to shortest paths and add to collection
-    for(const auto& state: algorithmResultStates) {
-        shortestPathList.emplace_back(state);
+    for(int i = 0; i < graphSize; ++i) {
+        shortestPathList.emplace_back(algorithmResultStates.getBestSolution(i));
     }
 
     //Add the (startingInstant, pathVector) pair at the appropriate node index
diff --git a/src/ShortestPath/Transit/TransitShortestPathContainer.h b/src/ShortestPath/Transit/TransitShortestPathContainer.h
index 923e816f9b1b5b3ee104c58b78d2dd7edfeb17b5..55166e4c813dee05d7e72396e07c24d81092369f 100644
--- a/src/ShortestPath/Transit/TransitShortestPathContainer.h
+++ b/src/ShortestPath/Transit/TransitShortestPathContainer.h
@@ -8,14 +8,17 @@
 
 #include <vector>
 #include "TransitShortestPath.h"
+#include "TransitStateContainer.h"
 
 class TransitShortestPathContainer {
 private:
     std::vector<std::vector<std::pair<int, std::vector<TransitShortestPath>>>> container; //NodeVector< PairVector<Pair<Instant, NodeVector<ShortestPath> >> >
 
 public:
+    explicit TransitShortestPathContainer(int size) { container.resize(size); }
+    explicit TransitShortestPathContainer(size_t size) { container.resize(size); }
     void addShortestPathCollection(int startNodeIndex, const std::pair<int, std::vector<TransitShortestPath>>& shortestPathList);
-    void addShortestPathCollection(int startNodeIndex, int startingInstant, int graphSize, const std::vector<TransitAlgorithmState>& algorithmResultStates);
+    void addShortestPathCollection(int startNodeIndex, int startingInstant, int graphSize, const TransitStateContainer& algorithmResultStates);
     std::vector<std::pair<int, std::vector<TransitShortestPath>>>::iterator getShortestPathsFromTime(int startNodeIndex, int earliestStartInstant);
     std::pair<int, TransitShortestPath> getShortestPathToYFromTime(int startNodeIndex, int earliestStartInstant, int goalNode);