fix beam_search

Author: tsutaj

docs/marathon/beam_search.md

@@ -1,3 +1,8 @@
+---
+title: ビームサーチ
+documentation_of: ./marathon/beam_search.cpp
+---
+
 - `BeamSearch` と `BeamSearchWithHash` の 2 種類が利用可能
     - 重複した状態を除去したいときは `BeamSearchWithHash` を使う。そうでないときは `BeamSearch` を使う。
     - `BeamSearch` だと $\log$ が落ちるはずなので、こっちのほうがイテレーションは回る
@@ -9,4 +14,4 @@
     - `to_next_state`: 軽量化した状態を受けとり、それを元の (軽量化していない) 状態に変換する関数
 
 AHC040 で使用実績あり。  
-https://atcoder.jp/contests/ahc040/submissions/61300845
+[https://atcoder.jp/contests/ahc040/submissions/61300845](https://atcoder.jp/contests/ahc040/submissions/61300845)

marathon/beam_search.cpp

@@ -1,8 +1,3 @@
-/**
- * @brief ビームサーチ
- * @docs docs/marathon/beam_search.md
- */
-
 template <typename BeamSearchState, typename LiteBeamSearchState>
 class BeamSearchBase {
    public:
@@ -15,7 +10,7 @@ class BeamSearchBase {
         const function<void(BeamSearchState &)> &add_next_lite_states,
         const function<BeamSearchState(LiteBeamSearchState &, BeamSearchState)> &to_next_state) = 0;
     const BeamSearchState &get_best_state() const {
-        if(states_.empty()) {
+        if (states_.empty()) {
             throw runtime_error("No states are registered.");
         }
         return *min_element(states_.begin(), states_.end());
@@ -45,14 +40,14 @@ class BeamSearch : public BeamSearchBase<BeamSearchState, LiteBeamSearchState> {
         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_++) {
+        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_]);
         }
         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++) {
+        for (int i = 0; i < num_select; i++) {
             next_states[i] = to_next_state(lite_states_[i], this->states_[lite_states_[i].state_idx]);
         }
         this->states_ = move(next_states);
@@ -71,7 +66,7 @@ class BeamSearchWithHash : public BeamSearchBase<BeamSearchState, LiteBeamSearch
         state.state_idx = this->current_state_idx_;
         lite_states_.emplace(state);
         // remove the worst (biggest) state.
-        while((int)lite_states_.size() > this->beam_width_) {
+        while ((int)lite_states_.size() > this->beam_width_) {
             lite_states_.pop();
         }
     }
@@ -82,17 +77,21 @@ class BeamSearchWithHash : public BeamSearchBase<BeamSearchState, LiteBeamSearch
         // 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_++) {
+        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_]);
         }
         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--) {
+            lite_states_vec[i] = lite_states_.top();
+            lite_states_.pop();
+        }
         vector<BeamSearchState> next_states(num_select);
         int num_next_states = 0;
-        while(lite_states_.size()) {
-            auto lite_state = lite_states_.top();
-            lite_states_.pop();
-            if(hash_values_.count(lite_state.hash_value)) {
+        for (int i = 0; i < (int)lite_states_vec.size(); i++) {
+            auto &lite_state = lite_states_vec[i];
+            if (hash_values_.count(lite_state.hash_value)) {
                 continue;
             }
             hash_values_.insert(lite_state.hash_value);
@@ -111,6 +110,7 @@ class BeamSearchWithHash : public BeamSearchBase<BeamSearchState, LiteBeamSearch
 template <typename score_t>
 struct BeamSearchStateBase {
     score_t score;
+    ll hash_value;
     BeamSearchStateBase() : score(0) {}
     virtual bool operator<(const BeamSearchStateBase &state) const = 0;
 };