diff --git a/download.h b/download.h
index e9d49d9..668c314 100644
--- a/download.h
+++ b/download.h
@@ -113,8 +113,6 @@ namespace Download {
             const std::vector<Action>& actions,
             size_t num_turns_to_replicate
     ) {
-
-    std::cout << "Initialising game with " << +num_players << " players, " << +num_suits << " suits and hand size " << +hand_size << std::endl;
       auto game = std::unique_ptr<Hanabi::HanabiStateIF>(new Hanabi::HanabiState<num_suits, num_players, hand_size>(deck));
       std::uint8_t index;
       for (size_t i = 0; i < num_turns_to_replicate; i++) {
diff --git a/game_state.hpp b/game_state.hpp
index 01b01bb..5379a6d 100644
--- a/game_state.hpp
+++ b/game_state.hpp
@@ -406,12 +406,11 @@ namespace Hanabi {
     #define RETURN_PROBABILITY                                        \
     if (_position_tablebase.contains(id_of_state)) {                  \
         ASSERT(_position_tablebase[id_of_state] == best_probability); \
-    } else {                                                          \
-        _position_tablebase[id_of_state] = best_probability;          \
-    }                                                                 \
+    }                                                          \
+    _position_tablebase[id_of_state] = best_probability;          \
+                                                                     \
     return best_probability;
 
-
     #define UPDATE_PROBABILITY(new_probability)                     \
     best_probability = std::max(best_probability, new_probability); \
     if (best_probability == 1) {                                    \
@@ -429,6 +428,9 @@ namespace Hanabi {
         if(_pace < 0 || _endgame_turns_left == 0) {
             return 0;
         }
+        if (_position_tablebase.contains(id_of_state)) {
+            return _position_tablebase[id_of_state];
+        }
 
         // TODO: Have some endgame analysis here?
 
@@ -500,8 +502,7 @@ namespace Hanabi {
             UPDATE_PROBABILITY(probability_stall);
         }
 
-        _position_tablebase[id_of_state] = best_probability;
-        return best_probability;
+        RETURN_PROBABILITY;
     }
 
     template<suit_t num_suits, player_t num_players, hand_index_t hand_size>
@@ -520,9 +521,18 @@ namespace Hanabi {
         id *= max_num_clues + 1;
         id += _num_clues;
 
-        // encode draw pile size
-        id *= _initial_draw_pile_size;
-        id += _weighted_draw_pile_size;
+        // we can encode draw pile size and extra turn in one metric, since we only have extra turns if draw pile is empty
+        const std::uint8_t draw_pile_size_and_extra_turns = [this]() -> uint8_t {
+            if(_endgame_turns_left == no_endgame) {
+                return _weighted_draw_pile_size + num_players;
+            }
+            else {
+                return _endgame_turns_left;
+            }
+        }();
+
+        id *= _initial_draw_pile_size + num_players;
+        id += draw_pile_size_and_extra_turns;
 
         // encode positions of cards that started in hands
         id = id << _num_useful_cards_in_starting_hands;