update search

tsutaj · tsutaj · commit a742005f37bd · 2025-06-25T11:24:55.000+09:00
diff --git a/docs/marathon/beam_search.md b/docs/marathon/beam_search.md
@@ -9,9 +9,15 @@ documentation_of: ./marathon/beam_search.cpp
 - 状態 `BeamSearchState` と、軽量化した状態 `LiteBeamSearchState` の両方を定義して使う
     - 軽量化した状態は、参照する `BeamSearchState` のインデックス +α の情報を持たせる
     - 両方ともスコア計算の実装が必須
-- 以下の 2 つの関数を定義してビームサーチを実行する
+- 以下の関数を定義する
     - `add_next_lite_states`: 状態を受けとり、次の状態を軽量化したものを返す関数
     - `to_next_state`: 軽量化した状態を受けとり、それを元の (軽量化していない) 状態に変換する関数
+- 概ね次のようにしてビームサーチをする
+  - 初期状態を登録
+  - ターン数だけ次を繰り返す
+    - 最初のターン以外: `reconstruct_states` で、軽量化した状態を元の状態に変換
+    - `add_lite_states` で、軽量化した状態を登録
+      - どのターンに対しても登録可能
 
-AHC040 で使用実績あり。  
-[https://atcoder.jp/contests/ahc040/submissions/61300845](https://atcoder.jp/contests/ahc040/submissions/61300845)
+AHC049 で使用実績あり。  
+[https://atcoder.jp/contests/ahc049/submissions/67053066](https://atcoder.jp/contests/ahc049/submissions/67053066)
diff --git a/marathon/beam_search.cpp b/marathon/beam_search.cpp
@@ -1,14 +1,15 @@
+#include <memory>
+
 template <typename BeamSearchState, typename LiteBeamSearchState>
 class BeamSearchBase {
    public:
     BeamSearchBase() = default;
     BeamSearchBase(int beam_width) : beam_width_(beam_width) {}
     void set_beam_width(int beam_width) { beam_width_ = beam_width; }
     void register_state(const BeamSearchState &state) { states_.emplace_back(state); }
-    virtual void register_lite_state(LiteBeamSearchState &state) = 0;
-    virtual void search(
-        const function<void(BeamSearchState &)> &add_next_lite_states,
-        const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) = 0;
+    virtual void register_lite_state(LiteBeamSearchState &lite_state, const BeamSearchState& state) = 0;
+    virtual void reconstruct_states(const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) = 0;
+    virtual void add_lite_states(const function<void(BeamSearchState &)> &add_next_lite_states) = 0;
     const BeamSearchState &get_best_state() const {
         if (states_.empty()) {
             throw runtime_error("No states are registered.");
@@ -32,26 +33,29 @@ template <class BeamSearchState, class LiteBeamSearchState>
 class BeamSearch : public BeamSearchBase<BeamSearchState, LiteBeamSearchState> {
    public:
     using BeamSearchBase<BeamSearchState, LiteBeamSearchState>::BeamSearchBase;
-    void register_lite_state(LiteBeamSearchState &state) override {
-        state.state_idx = this->current_state_idx_;
-        lite_states_.emplace_back(state);
+    void register_lite_state(LiteBeamSearchState &lite_state, const BeamSearchState& state) override {
+        lite_state.state_ref = make_shared<BeamSearchState>(state);
+        lite_states_.emplace_back(lite_state);
     }
-    void search(
-        const function<void(BeamSearchState &)> &add_next_lite_states,
-        const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) override {
-        lite_states_.clear();
-        for (this->current_state_idx_ = 0; this->current_state_idx_ < this->states_.size();
-             this->current_state_idx_++) {
-            add_next_lite_states(this->states_[this->current_state_idx_]);
-        }
+
+    void reconstruct_states(const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) override {
         const int num_select = min((int)lite_states_.size(), this->beam_width_);
         nth_element(lite_states_.begin(), lite_states_.begin() + num_select, lite_states_.end());
         vector<BeamSearchState> next_states(num_select);
         for (int i = 0; i < num_select; i++) {
-            next_states[i] = to_next_state(lite_states_[i], this->states_[lite_states_[i].state_idx]);
+            next_states[i] = to_next_state(lite_states_[i], *lite_states_[i].state_ref);
         }
         this->states_ = move(next_states);
     }
+
+    void add_lite_states(const function<void(BeamSearchState &)> &add_next_lite_states) override {
+        lite_states_.clear();
+        for (this->current_state_idx_ = 0; this->current_state_idx_ < this->states_.size();
+             this->current_state_idx_++) {
+            add_next_lite_states(this->states_[this->current_state_idx_]);
+        }
+    }
+
     int num_lite_states() const override { return lite_states_.size(); }
 
    private:
@@ -62,25 +66,12 @@ template <class BeamSearchState, class LiteBeamSearchState>
 class BeamSearchWithHash : public BeamSearchBase<BeamSearchState, LiteBeamSearchState> {
    public:
     using BeamSearchBase<BeamSearchState, LiteBeamSearchState>::BeamSearchBase;
-    void register_lite_state(LiteBeamSearchState &state) override {
-        state.state_idx = this->current_state_idx_;
-        lite_states_.emplace(state);
-        // remove the worst (biggest) state.
-        while ((int)lite_states_.size() > this->beam_width_) {
-            lite_states_.pop();
-        }
+    void register_lite_state(LiteBeamSearchState &lite_state, const BeamSearchState& state) override {
+        lite_state.state_ref = make_shared<BeamSearchState>(state);
+        lite_states_.emplace(lite_state);
     }
 
-    void search(
-        const function<void(BeamSearchState &)> &add_next_lite_states,
-        const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) override {
-        // NOTE: alternative of clear() because STL doesn't have priority_queue::clear().
-        lite_states_ = priority_queue<LiteBeamSearchState>();
-        hash_values_.clear();
-        for (this->current_state_idx_ = 0; this->current_state_idx_ < this->states_.size();
-             this->current_state_idx_++) {
-            add_next_lite_states(this->states_[this->current_state_idx_]);
-        }
+    void reconstruct_states(const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) override {
         const int num_select = min((int)lite_states_.size(), this->beam_width_);
         vector<LiteBeamSearchState> lite_states_vec(lite_states_.size());
         for (int i = (int)lite_states_vec.size() - 1; i >= 0; i--) {
@@ -89,37 +80,48 @@ class BeamSearchWithHash : public BeamSearchBase<BeamSearchState, LiteBeamSearch
         }
         vector<BeamSearchState> next_states(num_select);
         int num_next_states = 0;
-        for (int i = 0; i < (int)lite_states_vec.size(); i++) {
+        for (int i = 0; i < num_select; i++) {
             auto &lite_state = lite_states_vec[i];
             if (hash_values_.count(lite_state.hash_value)) {
                 continue;
             }
             hash_values_.insert(lite_state.hash_value);
-            next_states[num_next_states++] = to_next_state(lite_state, this->states_[lite_state.state_idx]);
+            next_states[num_next_states++] = to_next_state(lite_state, *lite_state.state_ref);
         }
         next_states.resize(num_next_states);
         this->states_ = move(next_states);
     }
+
+    void add_lite_states(const function<void(BeamSearchState &)> &add_next_lite_states) override {
+        // NOTE: alternative of clear() because STL doesn't have priority_queue::clear().
+        lite_states_ = priority_queue<LiteBeamSearchState>();
+        hash_values_.clear();
+        for (this->current_state_idx_ = 0; this->current_state_idx_ < this->states_.size();
+             this->current_state_idx_++) {
+            add_next_lite_states(this->states_[this->current_state_idx_]);
+        }
+    }
+
     int num_lite_states() const override { return lite_states_.size(); }
 
    private:
     priority_queue<LiteBeamSearchState> lite_states_;
     set<ll> hash_values_;
 };
 
-template <typename score_t>
-struct BeamSearchStateBase {
-    score_t score;
-    ll hash_value;
-    BeamSearchStateBase() : score(0) {}
-    virtual bool operator<(const BeamSearchStateBase &state) const = 0;
-};
+// template <typename score_t>
+// struct BeamSearchState {
+//     score_t score;
+//     ll hash_value;
+//     BeamSearchStateBase() : score(0), hash_value(0) {}
+//     ~BeamSearchStateBase() = default;
+// };
 
-template <typename score_t>
-struct LiteBeamSearchStateBase {
-    int state_idx;
-    score_t score;
-    ll hash_value;
-    LiteBeamSearchStateBase() : state_idx(-1), score(0), hash_value(0) {}
-    virtual bool operator<(const LiteBeamSearchStateBase &state) const = 0;
-};
+// template <typename score_t>
+// struct LiteBeamSearchState {
+//     shared_ptr<BeamSearchState<score_t>> state_ref;
+//     score_t score;
+//     ll hash_value;
+//     LiteBeamSearchState() : state_ref(nullptr), score(0), hash_value(0) {}
+//     ~LiteBeamSearchState() = default;
+// };
diff --git a/marathon/chokudai_search.cpp b/marathon/chokudai_search.cpp
@@ -9,7 +9,7 @@ class ChokudaiSearch {
         search_time_ = search_time;
         num_iter_ = 0;
         time_check_iter_ = time_check_iter;
-        max_num_states_.resize(max_turns, 1 << 30);
+        max_num_states_.resize(max_turns + 1, 1 << 30);
     }
     void set_max_num_states(int max_num_state) {
         fill(max_num_states_.begin(), max_num_states_.end(), max_num_state);
@@ -24,6 +24,13 @@ class ChokudaiSearch {
             states_[turns].pop_max();
         }
     }
+    size_t size(int turns) const {
+        return states_[turns].size();
+    }
+    const State& worst_state(int turns) const {
+        assert(states_[turns].size() > 0);
+        return states_[turns].top_max();
+    }
     void search(Timer &timer, const auto &add_next_states) {
         assert(num_iter_ >= 0);
         const double start_time = timer.getTime();
@@ -41,7 +48,12 @@ class ChokudaiSearch {
     TIME_OVER:;
         fprintf(stderr, "chokudai search: num_iter = %d\n", num_iter_);
     }
-    const State &get_best_state() const { return states_.back().top_min(); }
+    const State &get_best_state() const {
+        if (states_.back().empty()) {
+            throw runtime_error("No states are registered.");
+        }
+        return states_.back().top_min();
+    }
 
    private:
     double search_time_;
