diff --git a/include/game_state.h b/include/game_state.h
index 6fcf0c7..e1476eb 100644
--- a/include/game_state.h
+++ b/include/game_state.h
@@ -226,6 +226,12 @@ namespace Hanabi
 
     probability_t check_play_or_discard(hand_index_t index, bool play);
 
+    // For the current state, returns whether we will save it in the lookup table.
+    // By default, this is just constant true, but we might want to trade memory for speed, i.e.
+    // store less states, which will reduce memory consumption at the cost of re-computing some of the values
+    // when re-visiting the states.
+    bool save_state_to_map();
+
     static constexpr uint8_t no_endgame = std::numeric_limits<uint8_t>::max();
 
     // Usual game state
@@ -260,6 +266,7 @@ namespace Hanabi
     map_type<unsigned long, probability_t> _position_tablebase;
 
     std::uint64_t _enumerated_states{};
+
   };
 
   template<std::size_t num_suits, player_t num_players, std::size_t hand_size>
diff --git a/include/game_state.hpp b/include/game_state.hpp
index 06a992e..7d68268 100644
--- a/include/game_state.hpp
+++ b/include/game_state.hpp
@@ -982,6 +982,13 @@ namespace Hanabi
     return _actions_log.top().action_type;
   }
 
+  template<suit_t num_suits, player_t num_players, hand_index_t hand_size>
+  bool HanabiState<num_suits, num_players, hand_size>::save_state_to_map()
+  {
+    return true;
+    return _draw_pile.size() != 2;
+  }
+
   template<suit_t num_suits, player_t num_players, hand_index_t hand_size>
   probability_t HanabiState<num_suits, num_players, hand_size>::evaluate_state()
   {
@@ -1381,7 +1388,10 @@ namespace Hanabi
     {
       ASSERT(_position_tablebase[id] == probability);
     }
-    _position_tablebase[id] = probability;
+    if (save_state_to_map())
+    {
+      _position_tablebase[id] = probability;
+    }
   }
 
 } // namespace Hanabi
\ No newline at end of file