競技プログラミング学習

色々な問題について自分なりの方針などをまとめていきます。

ビームサーチをライブラリ化する【応用編】

前回の【基礎編】*1の記事では、ビームサーチの中でも一般に広く用いられていると思われる実装例を紹介しました。基本的にボトルネックとなるのは状態のコピーであり、コピーコストを軽くしたりスコアだけ先に計算してコピーを減らすなどして高速化を図りました。
今回紹介するのは、そもそも状態のコピーをやめてしまおうという実装方法です。一般にビームサーチにおけるノードは木構造を成していますから、毎回木を深さ優先探索(DFS)等で走査してスコアやハッシュ値を計算します。C++での実装例を3種類紹介しますが、いずれもポインタを駆使した実装になります。
本記事の内容はrhooさんによる解説記事 高速なビームサーチが欲しい!!! および 爆速ビームサーチライブラリを作る を参考にしていますので、こちらもぜひお読みください。
基礎編および応用編で紹介した実装はGitHub上でも公開しています。自由に使用していただいて構いませんが、結果に責任は負えませんことをご了承ください。

どんな問題で活躍するのか？

例えば前回題材とした「ゲーム実況者Xの挑戦」等のグリッド上を移動し続けるような問題では、グリッドの情報や各マスの訪問履歴などのコピーコストがかなり大きくなる場合があります。訪問履歴をビット列で保持する等コピーコストを小さくする工夫はできますが、それでもコストが大きそうな場合は今回紹介するコピーしない形式の実装も選択肢に入ります。
コピーコストが小さくなるように状態を設定できたら前回のような実装、それが難しいようなら今回のような実装が適しているでしょう。

実装する際の注意点

子ノードが全て不採用となった場合、そのノードはもはや不要となります。しかし不要となったノードを放置しておくと、その後のDFSにおいてもう調査する必要がない部分木を何度も行き来する可能性があります。そうなっては困るので、不採用になった葉から順に親をたどるなどして不要なノードは消去していく必要があります。これを可能にするために、各ノードには親および子ノードへのポインタをもたせることにします。

実装例１：生ポインタを用いる方法

実装方針

基礎編の実装例では各ノードに状態をコピーしていましたが、今回はノードは状態を保持しません。初期状態を一つ用意してDFSによって状態の更新と復元を繰り返しながら各ノードを調査していきます。以下のような形で実装してみます。

  1. 各ノードには親ノードへのポインタと子ノードへのポインタの配列を用意する。
  2. 初期状態を根として、子ノードへのポインタをたどりDFSで木を走査する。
  3. 子ノードへのポインタの配列が空ならばそのノードが葉であるはずなので、全ての操作を試して候補に追加する。
  4. 候補の中から上位beam_width個を採用し、葉から順に不要なノードを検出していく。
  5. 2. から 4. を最後まで繰り返す。

この大まかな方針は以降も変わりません。問題は不要なノードの検出方法とその扱いです。

実装例

ノード構造体に親ノードおよび子ノードへのポインタと、不要になったかどうかを表すexpiredという変数を用意します。

struct Node {
    Operation op;
    Node* parent;
    vector<Node*> children;
    bool expired;

    Node(Operation op, Node* parent);
};
Node::Node(Operation op, Node* parent) : op(op), parent(parent) {
    expired = false;
}

前回と同じように、スコアの比較は一時ノードで行うことにします。今回は本来のスコアraw_scoreと評価スコア(比較で使う値)eval_scoreをそれぞれ用意しました。

using ull = unsigned long long;

struct TemporaryNode {
    int raw_score, eval_score;
    ull hash;
    Operation op;
    Node* parent;

    TemporaryNode(int raw_score, int eval_score, ull hash, Operation op);
};
TemporaryNode::TemporaryNode(int raw_score, int eval_score, ull hash, Operation op) :
raw_score(raw_score), eval_score(eval_score), hash(hash), op(op) {
    parent = nullptr;
}

前回と同様に現在の状態を表すState構造体を用意します。ただし前回と違って状態を戻すという作業が加わるので、状態を更新した際に復元するための情報を返すようにします。

// 状態を復元するのに必要な情報をまとめた構造体
struct Restore {
    // ここに情報を書く
    Restore();
};

struct State {
    // (略)

    State();
    int score() const;
    ull hash() const;
    TemporaryNode try_move(const Operation& op) const;
    Restore apply_move(const Operation& op); // 状態を戻すための情報を返す

    void roll_back(const Restore& res, const Operation& op); // 状態を戻す関数が加わった
};

準備ができたのでDFSを実装します。現在の状態と根ノードへのポインタをもつTree構造体を用意し、その中にdfs関数を書きました。根ノードは操作や親ノードをもたないので、今回はそれぞれ-1とnullptrで初期化しました*2。

struct Tree {
    State state;
    Node* root_node;

    Tree(State& state);
    void dfs(Node* node_ptr, vector<TemporaryNode>& temp_nodes, bool single);
};
Tree::Tree(State& state) : state(state) {
    root_node = new Node(Operation{-1}, nullptr);
}
void Tree::dfs(Node* node_ptr, vector<TemporaryNode>& temp_nodes, bool single) {
    // 子ノードが存在しないなら葉なので、候補を追加して終了
    if(node_ptr->children.empty()) {
        for(const auto& op : valid_operations) {
            temp_nodes.emplace_back(state.try_move(op));
            temp_nodes.back().parent = node_ptr;
        }
        node_ptr->expired = true;
        return;
    }

    // 使われない子ノードを削除
    node_ptr->children.erase(remove_if(node_ptr->children.begin(), node_ptr->children.end(),
    [](Node* child_ptr) { return child_ptr->expired; }), node_ptr->children.end());

    // 一本道なら状態を戻さない
    bool next_single = single && ((int)node_ptr->children.size() == 1);

    auto node_backup = node_ptr;
    // 残った子ノードを走査
    for(auto& child_ptr : node_ptr->children) {
        Restore res = state.apply_move(child_ptr->op);
        dfs(child_ptr, temp_nodes, next_single);
        if(!next_single) state.roll_back(res, child_ptr->op);
    }
    if(!next_single) root_node = node_backup;

    // ノードを不要としておき、必要なら後で復活させる
    node_ptr->expired = true;
}

実装について二つ補足します。もし根から途中まで一本道だったとすれば、初めて分岐が発生するノードまで状態を戻せば十分で、毎回初期状態まで戻す必要はありません。それをsingleという変数で判定しています。すなわち、ここまでが一本道でありかつ子が一つしかないならnext_singleもtrueとなるようにしています。
もう一つの補足は、不要になったノードをどのように判定しているかです。本来であれば不要になったノードは葉から順にdelete演算子で削除したいのですが、削除すべきノードは親からも参照されていることに注意してください。言い換えると、削除されなかったノードの一部は子を参照していたダングリングポインタ*3を保持したままになります。これは次にDFSを実行した際に問題を起こしますから、ノードをdeleteするのは得策ではなさそうです。仕方ないのでメモリの効率化は諦めて、まだ必要なノードだけ復活させることにしました。すなわち、DFSをしながら一旦全てのノードを不要としておき、採用されたノードから親をたどって順次ノードを復活させていきます。
やや強引ですが、これでDFSを実装できました。とはいえ不要なノードを残しておく意味はないですし、効率的に削除できればノードを復活させるといった作業もしなくて済みます。後で紹介する2つの実装はこの問題に対する解決策となっています。
最後にビームサーチ本体の実装です。基本的には前回の実装をベースにしています。候補の追加はDFSに任せて、採用すると決定したノードから親をたどってノードを復活させます。

vector<Operation> BeamSearch(const int max_depth, const int beam_width) {
    State init_state;
    Tree tree(init_state);

    vector<TemporaryNode> final_nodes;

    unordered_set<ull> fields;
    vector<TemporaryNode> temp_nodes;

    for(int turn = 1; turn <= max_depth; turn++) {
        fields.clear();
        temp_nodes.clear();

        tree.dfs(tree.root_node, temp_nodes, true);
        // 最後のターンなら一時ノードの情報を保存して終了
        if(turn == max_depth) {
            final_nodes = temp_nodes;
            break;
        }

        int node_size = temp_nodes.size();
        // 候補がビーム幅より多いなら上位beam_width個を選ぶ
        if(node_size > beam_width) {
            nth_element(temp_nodes.begin(), temp_nodes.begin() + beam_width, temp_nodes.end(),
            [](TemporaryNode& n1, TemporaryNode& n2) {
                return n1.eval_score > n2.eval_score;
            });
        }
        // 仮ノードの情報から実際にノードを更新する
        for(int i = 0; i < min(beam_width, node_size); i++) {
            if(fields.count(temp_nodes[i].hash)) continue;
            fields.insert(temp_nodes[i].hash);
            temp_nodes[i].parent->children.emplace_back(
                new Node(temp_nodes[i].op, temp_nodes[i].parent)
            );

            // 採用されたノードから親をたどりノードを復活させる
            Node* node_ptr = temp_nodes[i].parent;
            while(node_ptr && node_ptr->expired) {
                node_ptr->expired = false;
                node_ptr = node_ptr->parent;
            }
        }
    }

    // 最良の状態を選択
    int arg_best = -1;
    int best_score = 0;
    for(int i = 0; i < (int)final_nodes.size(); i++) {
        if(final_nodes[i].raw_score > best_score) {
            arg_best = i;
            best_score = final_nodes[i].raw_score;
        }
    }
    assert(arg_best != -1);
    Operation op = final_nodes[arg_best].op;
    auto ptr = final_nodes[arg_best].parent;

    vector<Operation> result{op};
    // 操作の復元
    while(ptr->parent) {
        result.emplace_back(ptr->op);
        ptr = ptr->parent;
    }
    reverse(result.begin(), result.end());
    return result;
}

実装例２：スマートポインタを用いる方法

スマートポインタを用いる利点

実装例１では、不正なポインタが残ってしまわないようにノードの消去をしませんでした。しかしスマートポインタを用いるとこの問題を巧妙に回避できます。今回は以下の二種類のポインタを利用します。

  • shared_ptr ・・・同じ参照先を複数のポインタで共有できる。最後のshared_ptrが破棄される時に参照先を自動で解放する。
  • weak_ptr ・・・shared_ptrの参照先のみを参照できる。weak_ptrが残っていても最後のshared_ptrは構わず参照先を解放する。

この性質の違いが重要です。そしてweak_ptrには参照先が解放されているかを返すexpiredという関数があるため、これを使えば不正なアクセスを避けることができます。
これをどう用いるのかと言うと、子ノードへのポインタはweak_ptrとし、逆に親ノードへのポインタはshared_ptrとします。各ノードは子からのshared_ptrによる参照があるため生きていますが、子が存在しなくなると親からのweak_ptrしか残らないため自動的に解放されます。挙動としては実装例１で述べたdelete演算子を用いたノードの消去に似ていますが、今回はweak_ptrを用いているため参照先が解放されているか簡単に判定できます。すなわち、次のDFSにおいて参照先が解放されているweak_ptrを削除するだけで済みます。

実装例

ノード構造体は以下のようになります。

struct Node {
    Operation op;
    shared_ptr<Node> parent;
    vector<weak_ptr<Node>> children;

    Node(Operation op, shared_ptr<Node> parent);
};
Node::Node(Operation op, shared_ptr<Node> parent) : op(op), parent(parent) {}

一時ノードにおける親ノードへのポインタも同様にshared_ptrに変更します。DFSの実装も、ノードのexpired変数をtrueに変更していた行を削除するだけで後は同じです。それではビームサーチ本体の実装を見てみます。

vector<Operation> BeamSearch(const int max_depth, const int beam_width) {
    State init_state;
    Tree tree(init_state);

    vector<shared_ptr<Node>> current_nodes;
    vector<TemporaryNode> final_nodes;

    unordered_set<ull> fields;
    vector<TemporaryNode> temp_nodes;

    for(int turn = 1; turn <= max_depth; turn++) {
        fields.clear();
        temp_nodes.clear();

        tree.dfs(tree.root_node, temp_nodes, true);
        // 最後のターンなら一時ノードの情報を保存して終了
        if(turn == max_depth) {
            final_nodes = temp_nodes;
            break;
        }

        int node_size = temp_nodes.size();
        // 候補がビーム幅より多いなら上位beam_width個を選ぶ
        if(node_size > beam_width) {
            nth_element(temp_nodes.begin(), temp_nodes.begin() + beam_width, temp_nodes.end(),
            [](TemporaryNode& n1, TemporaryNode& n2) {
                return n1.eval_score > n2.eval_score;
            });
        }
        // 仮ノードの情報から実際にノードを更新する
        current_nodes.clear();
        for(int i = 0; i < min(beam_width, node_size); i++) {
            if(fields.count(temp_nodes[i].hash)) continue;
            fields.insert(temp_nodes[i].hash);
            current_nodes.emplace_back(make_shared<Node>(temp_nodes[i].op, temp_nodes[i].parent));
            temp_nodes[i].parent->children.emplace_back(current_nodes.back());
        }
    }
    // (略)
}

実装例１とほとんど同じです。採用された葉のノードが、shared_ptrの配列current_nodesに保管されています。こうすることで一時ノードを保管しているtemp_nodesの中身が消去された時点で、不採用のノードはshared_ptrからの参照がなくなり自動的に解放されます。そして子が全て解放されたノードも自動的に解放される、といった要領で連鎖的に不要なノードが削除されていきます。

実装例３：二重連鎖木による高速化

これまでの実装の欠点と二重連鎖木を用いる利点

実装例１では不要なノードが残るため、メモリ使用量が大きくなったり必要なノードを復活させなければならないという欠点がありました。スマートポインタを用いた実装例２はこの問題点を解消できていますが、スマートポインタは丁寧にメモリを管理する分どうしても実行速度が下がってしまいます。生のポインタを使いながら不要なノードを消去する方法を考えたくなります。
実装例１では各ノードが全ての子ノードへのポインタを保持していたため、既に解放された子ノードへアクセスする危険がありました。これを回避するために、代表となる子ノードを設定しその子ノードへのポインタのみを保持することにします。そして各ノードには親ノードと子ノードに加えて前後の兄弟ノードへのポインタも用意します。言い換えると、今まで子へのポインタを並列に保持していたのを直列に変更します。このデータ構造を二重連鎖木と呼びます。
こうすると、不正なポインタが残ることを防げます。例えば削除するノードに前の兄弟がいないならそれは子の代表ですから、次の兄弟に代表を譲ってから削除します。もし前の兄弟も次の兄弟もいないなら自身が最後の子ですから、削除する時に親ノードも一緒に削除します。

実装例

ノード構造体は以下のようになります。高速化のためビームサーチの実装方法を少し変えたので、それに合わせてinitializeという初期化関数を用意しました。

struct Node {
    Operation op;
    Node* parent;
    Node* child;
    Node* prev_sibling;
    Node* next_sibling;

    Node();
    void initialize(Operation op, Node* parent);
};
Node::Node() {
    op = -1;
    parent = nullptr;
    child = nullptr;
    prev_sibling = nullptr;
    next_sibling = nullptr;
}
void Node::initialize(Operation op, Node* parent) {
    this->op = op;
    this->parent = parent;
    child = nullptr;
    prev_sibling = nullptr;
    next_sibling = nullptr;
}

一時ノードやDFSは今までと同様に実装できます。ビームサーチ本体は以下のようにしました。

vector<Operation> BeamSearch(const int max_depth, const int beam_width) {
    State init_state;
    Tree tree(init_state);

    vector<TemporaryNode> final_nodes, temp_nodes;
    vector<Node*> current_nodes, next_nodes;

    constexpr int max_nodes = 15000; // 要調整
    vector<Node> valid_nodes(max_nodes);
    tree.root_node = &valid_nodes[0];
    vector<Node*> vacant;
    for(int i = max_nodes - 1; i > 0; i--) {
        vacant.emplace_back(&valid_nodes[i]);
    }

    unordered_set<uint> fields;

    for(int turn = 1; turn <= max_depth; turn++) {
        // (略)

        // 仮ノードの情報から実際にノードを更新する
        for(int i = 0; i < min(beam_width, node_size); i++) {
            if(fields.count(temp_nodes[i].hash)) continue;
            fields.insert(temp_nodes[i].hash);

            Node* parent = temp_nodes[i].parent;
            /*
            // 必要ならばノード数が足りているか確認する
            if(vacant.empty()) {
                std::cerr << "max_nodes を大きくしてください" << std::endl;
            }
            assert(!vacant.empty());
            */
            // 既に子がいるなら新たに代表の子として挿入する
            if(parent->child) {
                parent->child->prev_sibling = vacant.back();
                vacant.back()->initialize(temp_nodes[i].op, parent);

                parent->child->prev_sibling->next_sibling = parent->child;
                parent->child = parent->child->prev_sibling;
            }
            // 子がいないなら代表の子とする
            else {
                parent->child = vacant.back();
                vacant.back()->initialize(temp_nodes[i].op, parent);
            }
            next_nodes.emplace_back(vacant.back());
            vacant.pop_back();
        }
        // 子がいないノードを再帰的に削除する
        for(auto ptr : current_nodes) {
            while(!ptr->child) {
                // 前も後ろも兄弟がいる場合
                if(ptr->prev_sibling && ptr->next_sibling) {
                    ptr->prev_sibling->next_sibling = ptr->next_sibling;
                    ptr->next_sibling->prev_sibling = ptr->prev_sibling;
                    vacant.emplace_back(ptr);
                    break;
                }
                // 前の兄弟だけいる場合
                else if(ptr->prev_sibling && !ptr->next_sibling) {
                    ptr->prev_sibling->next_sibling = nullptr;
                    vacant.emplace_back(ptr);
                    break;
                }
                // 後ろの兄弟だけいる場合
                else if(!ptr->prev_sibling && ptr->next_sibling) {
                    ptr->next_sibling->prev_sibling = nullptr;
                    ptr->parent->child = ptr->next_sibling;
                    vacant.emplace_back(ptr);
                    break;
                }
                // 両方いない場合はさらに親も削除する
                else {
                    if(!ptr->parent) {
                        vacant.emplace_back(ptr);
                        break;
                    }
                    else {
                        vacant.emplace_back(ptr);
                        ptr = ptr->parent;
                        ptr->child = nullptr;
                    }
                }
            }
        }

        swap(current_nodes, next_nodes);
        next_nodes.clear();
    }
    // (略)
}

必要になるたびにnew演算子でノードを生成するのではなく、あらかじめ一定数のノードを用意しておき、不要になったノードは後で上書きするという実装にしました*4。

まとめ

今回はノードに状態をコピーしないビームサーチの実装方法を紹介しました。状態のコピーコストが大きくなってしまう場合はこちらの方が良い場合も多いかもしれません。
現在のところ、今回紹介した実装例の中では最後の二重連鎖木を用いた実装が優秀で一番ビーム幅を確保できるようです。ですが他にも良い実装方法があるかもしれませんので、皆さんもぜひ考えてみてください。実装例もあくまで一例ですので、内容を理解できた方は自分が使いやすいように一から実装してみるのも面白いと思います。

*1:本記事の内容は基礎編の内容がベースになっているために応用編としており、基礎編が簡単で応用編が難しいといった意図ではありません。

*2:より丁寧に実装するなら、操作の無効値を-1とするのではなく、std::optionalを用いてnulloptを代入することもできます。

*3:参照先のメモリが解放される等によって無効なメモリ領域を指しているポインタのことです。本来こうしたポインタにはnullptrを代入すべきですが、親ノードが保持しているポインタのうちどれが自身への参照なのかを判断する情報を子ノード側は保持していません。かといってその度に親ノードがもつ配列を走査していては時間がかかってしまいます。

*4:高速化のために色々と試している中、毎回new演算子でノードを生成すると少し遅い可能性があるとrhooさんから指摘をいただきました。ありがとうございます。

ビームサーチをライブラリ化する【基礎編】

7/30 追記：一部の実装を変更して高速化しました。(ビームサーチ内で毎回配列を定義するコストが無視できなかったため、一度だけ定義して毎回中身を全て消去する形にしました)



ヒューリスティックコンテストでよく用いられるビームサーチという手法があります。山登り/焼きなまし法と並んで必ずと言っていいほど紹介されるビームサーチですが、実装が大変と感じている方も多いと思います。
この記事ではビームサーチを既に知っている方向けに、ライブラリ化の指針を提供したいと考えています。ビームサーチ自体については簡単に説明する程度ですのでご了承ください*1。
今回は基礎編と題して、一般に広く用いられていると思われる実装方法を紹介します。応用編*2ではrhooさんの記事などで解説されている他の実装方法を紹介したいと考えています。本記事中のソースコードは自由に利用していただいて構いませんが、万一何らかの不利益が生じても筆者は責任を負えません。

説明にあたり具体的な例があるとわかりやすいと思うので、ビームサーチの練習問題としてよく取り上げられている「ゲーム実況者Xの挑戦」を題材にしてみたいと思います。
atcoder.jp
問題の概要としては、$N$ 個あるグリッド状のマップから $K$ 個を選び、全てのマップ上で同じ動き方(上下左右のいずれか)をしてなるべく多くのコインを取得することを目指します。時系列が重要なターン制の問題なのでビームサーチが有効そうな気がしますね。

ビームサーチの挙動を考える

全ての行動パターンを探索できれば厳密な最適解が得られるわけですが、毎ターン4つの選択肢があって2500ターンも続くので当然それは不可能です。そこでビームサーチでは、一部の評価が高い状態だけを残して次に進むことを繰り返すのでした*3。この挙動はどのような問題にも共通していますから、うまくライブラリ化できないか考えてみます。

実装例

実装の方針

ビームサーチで用いる各状態のことをノードと呼ぶことにします。上述の挙動を素直に実装するなら以下のようになるでしょうか。

  1. ノード構造体を定義する。残すノードを決められるように評価関数を用意する。
  2. 初期ノードを優先度付きキューに入れる。
  3. 優先度付きキューから評価が高い一定個数を取り出し、各ノードについて全ての遷移(今回は4通り)を試して次の優先度付きキューに全て入れる。
  4. 3. を最後まで繰り返す。

既に通過したマスで再度コインを取得してしまわないように、各マップについてどのマスを訪問済みかの情報をビット列で保持します。ではC++での実装例を見てみましょう。

ゲームの進行状況を表す構造体を用意する

まずは現在のゲームの進行状況を表すState構造体を定義し、それを各ノードにもたせることにします。その際、操作を表すOperation構造体やプレイヤーの情報をもつPlayer構造体を用意しています。
State構造体はビームサーチでコピーする情報を保持するので、なるべくコピーコストが小さくなるようにデータの持ち方を工夫したいです。ここが問題に固有の部分で、他はほとんど使い回せるような実装にします。

using ull = unsigned long long;

struct State {
    int now_turn;
    int coins;
    int penalty; // 罠を踏んだ時にペナルティを加算する
    ull zobrist_hash; // 盤面の重複を削除するためのハッシュ値

    vector<Player> players;
    vector<bitset<2500>> visited;
    string move_history;

    State();
    int score() const;
    ull hash() const;
    void apply_move(const Operation& op);
};
State::State() {
    // 初期化(初期盤面のハッシュ値の計算、プレイヤーの位置情報の取得など)
}
int State::score() const {
    // 評価関数を設定する
}
ull State::hash() const {
    // ハッシュ値、今回はzobrist hash
}
void State::apply_move(const Operation& op) {
    // ここにターンを進める処理を書く
    // (スコアとハッシュ値の計算、プレイヤーや訪問済み情報の更新など)
}

これまでの移動の履歴をmove_historyという文字列で保持して、最後に操作を復元できるようにします。

ノード構造体を用意する

続いてビームサーチで用いるノード構造体を定義します。

struct Node {
    State state;

    Node(State& state);
    int get_score() const;
    ull get_hash() const;
    void advance(const Operation& op);
    bool operator< (const Node& node) const;
};
Node::Node(State& state) : state(state) {}
int Node::get_score() const {
    return state.score();
}
ull Node::get_hash() const {
    return state.hash();
}
void Node::advance(const Operation& op) {
    state.apply_move(op);
}
bool Node::operator< (const Node& node) const {
    return get_score() < node.get_score();
}

ノードをコピーしてadvance関数を呼ぶことで状態が更新されます。ノードの実装はどの問題でも基本的に変更する必要がないようにしています。

ビームサーチを実装する

準備ができたので、いよいよビームサーチの実装です。各ノードからの遷移を書きやすくするために、ありうる操作を列挙したvalid_operationsというOperation型の配列を用意しています。

Node BeamSearch(State& init_state, const int max_depth, const int beam_width) {
    priority_queue<Node> nodes;
    nodes.emplace(init_state); // 初期状態を優先度付きキューに入れる
    unordered_set<ull> fields; // 重複除去用

    for(int turn = 1; turn <= max_depth; turn++) {
        priority_queue<Node> next_nodes;
        // 上位beam_width個だけ取り出す
        for(int i = 0; i < beam_width; i++) {
            if(nodes.empty()) break;
            Node node = nodes.top(); nodes.pop();
            // 可能な全ての遷移を試す
            for(auto& op : valid_operations) {
                Node next_node = node;
                next_node.advance(op);
                if(fields.count(next_node.get_hash())) continue; // 過去との重複をチェック
                next_nodes.emplace(next_node);
            }
        }

        while(!nodes.empty()) nodes.pop();
        // 重複を除去しながら上位beam_width個を採用する
        while(!next_nodes.empty() && (int)nodes.size() < beam_width) {
            if(fields.count(next_nodes.top().get_hash())) { // このターン内の重複をチェック
                next_nodes.pop();
                continue;
            }
            fields.insert(next_nodes.top().get_hash());
            nodes.emplace(next_nodes.top());
            next_nodes.pop();
        }
    }

    return nodes.top();
}

ビームサーチ自体もほとんど変更を加えず使い回せると思います。前述のように、問題に固有なのはコピーすべき状態をもつState構造体や各操作を表すOperation構造体になります。

この問題は実行制限時間が4秒で、ここまでの実装でビーム幅を200程度は確保できると思います。ですが実は無駄な処理をしていて高速化の余地が残っています。

高速化するには

スコアだけ先に計算して上位を決定する

ビームサーチで実行時間のボトルネックになるのは、多くの場合ノードのコピーでしょう。上の実装では毎ターン最大で beam_width * 操作の種類数 だけコピーが発生しますが、よく考えると結局採用されないノードが多数あり無駄なコピーが発生しています。
これを回避するには、ノードをコピーせずにスコア(とハッシュ値)だけ先に計算して上位を決定すれば良いです。今までState構造体には実際に操作をして状態を更新するapply_moveという関数しかありませんでしたが、状態を更新せずにスコアだけ計算するtry_moveという関数を用意します。

struct State {
    // ...他の部分は変更なし

    pair<int,ull> try_move(const Operation& op) const; // 状態を更新しないのでconstをつけておく
}
pair<int,ull> State::try_move(const Operation& op) const {
    // 一手進めた場合のスコアとハッシュ値を返す、ただし状態を更新はしない
}

ノード構造体にも対応するcalculateという関数を用意します。

struct Node {
    // ...他の部分は変更なし

    pair<int,ull> calculate(const Operation& op) const;
}
pair<int,ull> Node::calculate(const Operation& op) const {
    return state.try_move(op);
}

そしてビームサーチではノードで比較するのではなく、スコアなどの情報を保持した一時ノードを用意して比較します。盤面の情報をもたないのでコピーコストが軽くなっています。

random_device rnd;
mt19937 engine(rnd());
uniform_real_distribution<> randR(0.0, 1.0);
// スコアだけ計算して上位を選ぶために用いる仮ノード
struct TemporaryNode {
    int score;
    ull hash;
    int node_index;
    Operation op;
    double rand; // タイブレーク用

    TemporaryNode(int score, ull hash, int node_index, Operation& op);
};
TemporaryNode::TemporaryNode(int score, ull hash, int node_index, Operation& op) :
score(score), hash(hash), node_index(node_index), op(op) {
    rand = randR(engine);
}

スコアが同じ状態に差をつけるために乱数を入れておきました。
続いてビームサーチ本体です。先ほどは優先度付きキューで上位を取り出しましたが、C++には指定した個数だけ配列の上位を並び替えてくれるnth_elementという便利な関数があるのでそれを活用してみます。

Node BeamSearch(State& init_state, const int max_depth, const int beam_width) {
    vector<Node> nodes, next_nodes;
    nodes.emplace_back(init_state);
    nodes.back().move_history = Stack{nullptr};

    vector<TemporaryNode> temp_nodes; // スコア比較用の仮ノードを保管
    unordered_set<ull> fields; // 重複除去用

    for(int turn = 1; turn <= max_depth; turn++) {
        temp_nodes.clear();
        fields.clear();
        
        for(int i = 0; i < (int)nodes.size(); i++) {
            // 可能な全ての遷移を試す
            for(auto& op : valid_operations) {
                auto [next_score, next_hash] = nodes[i].calculate(op);
                temp_nodes.emplace_back(next_score, next_hash, i, op);

                // 必要なら重複除去
                if(fields.count(temp_nodes.back().hash)) {
                    temp_nodes.pop_back();
                }
                else {
                    fields.insert(temp_nodes.back().hash);
                }
            }
        }

        int node_size = temp_nodes.size();
        // 候補がビーム幅より多いなら上位beam_width個を選ぶ
        if(node_size > beam_width) {
            nth_element(temp_nodes.begin(), temp_nodes.begin() + beam_width, temp_nodes.end(),
            [](TemporaryNode& n1, TemporaryNode& n2) {
                if(n1.score == n2.score) {
                    return n1.rand > n2.rand;
                }
                return n1.score > n2.score;
            });
        }

        // 仮ノードの情報から実際にノードを更新する
        for(int i = 0; i < min(beam_width, node_size); i++) {
            int index = temp_nodes[i].node_index;
            next_nodes.emplace_back(nodes[index]);
            next_nodes.back().advance(temp_nodes[i].op);
            // 必要ならスコアとハッシュ値を確認
            // assert(next_nodes.back().score == temp_nodes[i].score);
            // assert(next_nodes.back().hash == temp_nodes[i].hash);
        }

        swap(nodes, next_nodes);
        next_nodes.clear();
    }

    int arg_best = -1, best_score = 0;
    for(int i = 0; i < (int)nodes.size(); i++) {
        if(nodes[i].get_score() > best_score) {
            arg_best = i;
            best_score = nodes[i].get_score();
        }
    }
    return nodes[arg_best];
}

これもほとんど変更せずに使い回せると思います。一時ノードには比較関数を定義していませんが、ラムダ式でnth_elementに直接渡しています。これで毎ターン最大beam_width個のノードのコピーで済むようになりました。

発展：永続スタックを用いた行動履歴の復元

記事のタイトルに基礎編と書きましたが、最後に一つ発展的な内容を紹介します。
今までは移動の履歴を文字列で保持していましたが、終盤になると文字列のサイズが大きくなりコピーコストも無視できません。そこで文字列を毎回コピーして直接保持するのではなく効率的に復元する方法を考えます。
各ノードに親ノードへのポインタをもたせておきたくなりますが、過去のノードは破棄されてしまいますし、かといって全て保存しておいたらメモリが大変なことになります。そこで永続データ構造と呼ばれるものを利用して履歴を保存してみます。詳しくはnoshi91さんの記事をご覧ください*4。ここではスマートポインタを用いてメモリを管理する実装方法を選択しました。もう参照されることがない履歴は自動的にメモリを解放してくれるので便利です。

// 行動を復元する永続stack
struct History {
    Operation op;
    shared_ptr<History> parent;
    History(const Operation& op, shared_ptr<History>& parent) : op(op), parent(parent) {}
};
struct Stack {
    shared_ptr<History> head;

    Operation top();
    Stack push(const Operation& op);
    Stack pop();
};
Operation Stack::top() {
    return head->op;
}
Stack Stack::push(const Operation& op) {
    return Stack({make_shared<History>(op, head)});
}
Stack Stack::pop() {
    return Stack({head->parent});
}

ノードには新たにスタックをもたせます。

struct Node {
    State state;
    Stack move_history; // これが追加された

    Node(State& state);
    int get_score() const;
    ull get_hash() const;
    pair<int,ull> calculate(const Operation& op) const;
    void advance(const Operation& op);
};

ビームサーチでは各ノードでスタックに操作を追加します。

Node BeamSearch(State& init_state, const int max_depth, const int beam_width) {
    vector<Node> nodes, next_nodes;
    nodes.emplace_back(init_state);
    nodes.back().move_history = Stack{nullptr};

    vector<TemporaryNode> temp_nodes; // スコア比較用の仮ノードを保管
    unordered_set<ull> fields; // 重複除去用

    for(int turn = 1; turn <= max_depth; turn++) {
        // (略)

        // 仮ノードの情報から実際にノードを更新する
        vector<Node> next_nodes;
        for(int i = 0; i < min(beam_width, node_size); i++) {
            int index = temp_nodes[i].node_index;
            next_nodes.emplace_back(nodes[index]);
            next_nodes.back().advance(temp_nodes[i].op);
            // 必要ならスコアとハッシュ値を確認
            // assert(next_nodes.back().score == temp_nodes[i].score);
            // assert(next_nodes.back().hash == temp_nodes[i].hash);
            fields.insert(next_nodes.back().get_hash());

            // 親ノードのスタックに操作を追加して新しいノードのスタックを作成する
            next_nodes.back().move_history = nodes[index].move_history.push(temp_nodes[i].op);
        }

        swap(nodes, next_nodes);
        next_nodes.clear();
    }

    // (略)
}

実際に操作を復元する際は以下のように用います。

State init_state;
constexpr int max_depth = 2500, beam_width = 1000;
Node result = BeamSearch(init_state, max_depth, beam_width);

vector<Operation> moves;
Stack move_history = result.move_history;
while(move_history.head) {
    Operation op = move_history.top();
    moves.emplace_back(op);
    move_history = move_history.pop();
}
reverse(moves.begin(), moves.end());

最終的な実装例

以上をまとめて次のようなライブラリになりました。この実装で今回の問題ではビーム幅1000以上を達成できています。まだまだ改良の余地はあると思いますので、皆さんも最強のビームサーチライブラリを自作しちゃいましょう。長くなってしまいましたが、最後までお読みいただきありがとうございました。

random_device rnd;
mt19937 engine(rnd());
uniform_real_distribution<> randR(0.0, 1.0);
using ull = unsigned long long;

struct State {
    // コピーすべき情報をここに書く

    State();
    int score() const;
    ull hash() const;
    pair<int,ull> try_move(const Operation& op) const;
    void apply_move(const Operation& op);
};
State::State() {

}
int State::score() const {

}
ull State::hash() const {

}
// 一手進めた場合のスコアとハッシュ値を返す、更新はしない
pair<int,ull> State::try_move(const Operation& op) const {

}
// 更新する
void State::apply_move(const Operation& op) {

}

// 行動を復元する永続stack
struct History {
    Operation op;
    shared_ptr<History> parent;
    History(const Operation& op, shared_ptr<History>& parent) : op(op), parent(parent) {}
};
struct Stack {
    shared_ptr<History> head;

    Operation top();
    Stack push(const Operation& op);
    Stack pop();
};
Operation Stack::top() {
    return head->op;
}
Stack Stack::push(const Operation& op) {
    return Stack({make_shared<History>(op, head)});
}
Stack Stack::pop() {
    return Stack({head->parent});
}

struct Node {
    State state;
    Stack move_history;

    Node(State& state);
    int get_score() const;
    ull get_hash() const;
    pair<int,ull> calculate(const Operation& op) const;
    void advance(const Operation& op);
};
Node::Node(State& state) : state(state) {}
int Node::get_score() const {
    return state.score();
}
ull Node::get_hash() const {
    return state.hash();
}
pair<int,ull> Node::calculate(const Operation& op) const {
    return state.try_move(op);
}
void Node::advance(const Operation& op) {
    state.apply_move(op);
}

// スコアだけ計算して上位を選ぶために用いる仮ノード
struct TemporaryNode {
    int score;
    ull hash;
    int node_index;
    Operation op;
    double rand; // タイブレーク用

    TemporaryNode(int score, ull hash, int node_index, Operation& op);
};
TemporaryNode::TemporaryNode(int score, ull hash, int node_index, Operation& op) :
score(score), hash(hash), node_index(node_index), op(op) {
    rand = randR(engine);
}

Node BeamSearch(State& init_state, const int max_depth, const int beam_width) {
    vector<Node> nodes, next_nodes;
    nodes.emplace_back(init_state);
    nodes.back().move_history = Stack{nullptr};

    vector<TemporaryNode> temp_nodes; // スコア比較用の仮ノードを保管
    unordered_set<ull> fields; // 重複除去用

    for(int turn = 1; turn <= max_depth; turn++) {
        temp_nodes.clear();
        fields.clear();
        
        for(int i = 0; i < (int)nodes.size(); i++) {
            // 可能な全ての遷移を試す
            for(auto& op : valid_operations) {
                auto [next_score, next_hash] = nodes[i].calculate(op);
                temp_nodes.emplace_back(next_score, next_hash, i, op);

                // 必要なら重複除去
                if(fields.count(temp_nodes.back().hash)) {
                    temp_nodes.pop_back();
                }
                else {
                    fields.insert(temp_nodes.back().hash);
                }
            }
        }

        int node_size = temp_nodes.size();
        // 候補がビーム幅より多いなら上位beam_width個を選ぶ
        if(node_size > beam_width) {
            nth_element(temp_nodes.begin(), temp_nodes.begin() + beam_width, temp_nodes.end(),
            [](TemporaryNode& n1, TemporaryNode& n2) {
                if(n1.score == n2.score) {
                    return n1.rand > n2.rand;
                }
                return n1.score > n2.score;
            });
        }

        // 仮ノードの情報から実際にノードを更新する
        for(int i = 0; i < min(beam_width, node_size); i++) {
            int index = temp_nodes[i].node_index;
            next_nodes.emplace_back(nodes[index]);
            next_nodes.back().advance(temp_nodes[i].op);
            // 必要ならスコアとハッシュ値を確認
            // assert(next_nodes.back().score == temp_nodes[i].score);
            // assert(next_nodes.back().hash == temp_nodes[i].hash);

            // 親ノードのスタックに操作を追加して新しいノードのスタックを作成する
            next_nodes.back().move_history = nodes[index].move_history.push(temp_nodes[i].op);
        }

        swap(nodes, next_nodes);
        next_nodes.clear();
    }

    int arg_best = -1, best_score = 0;
    for(int i = 0; i < (int)nodes.size(); i++) {
        if(nodes[i].get_score() > best_score) {
            arg_best = i;
            best_score = nodes[i].get_score();
        }
    }
    return nodes[arg_best];
}

*1:ヒューリスティックで活用できる手法を体系的に学びたい方は、最近発売されたthunder本「ゲームで学ぶ探索アルゴリズム実践入門」がおすすめです。ビームサーチのC++での実装例も提供されています。

*2:執筆できるか定かではないですが頑張りたいと思います。

*3:ただし現在のスコアが高ければ最終スコアも高いとは限らないので、現在の状況をどう評価するかは腕の見せ所です。

*4:永続データ構造を用いてビームサーチの行動を復元できることはてりーさんに教えてもらいました。ありがとうございます。

yukicoder score contest 6 writer記

少し時間が経ってしまいましたが、先日 yukicoder score contest 6 を開催させていただきました。今回はテスターは yunix さんにお願いしました。
優勝者は terry_u16 さんでした、おめでとうございます！

問題概要

縦スクロール型のシューティングゲームをプレイするプログラムを作成してください。
敵機は最上段に出現してグリッドを下に降りてくるので、最下段の自機を左右に動かしてうまく敵機を避けながら攻撃してね。
yukicoder.me

今回はリアクティブ形式の問題でした。毎ターン敵機がどの列に出現したか与えられるので、それに応じて自機をどう動かすかを決定します。
途中で敵機に衝突せずにゲームを終えられれば、シンプルな解法でもそこそこの点数が取れます。ただし衝突判定は結構バグらせやすく、短期コンテストだと厄介だったかもしれません。

想定解 (writer解/tester解)

序盤は基本的にプレイヤーの強化(=パワーの取得)を優先し、終盤で一気にスコアを稼ぐと良いというのは writer/tester 共通の見解でした。
この方針に則ったシンプルな貪欲法も結構強いのですが、ある程度シミュレーションで先読みして行動を決めた方がより強そうです。試しに使ってみたかったこともあってwriter解はモンテカルロ木探索(MCTS)のような何か*1、tester解はビームサーチでした。結論としてはこの問題はビームサーチの方が優秀だったように思います。

ビジュアライザ

前回に続き Siv3D というC++用のフレームワークを用いてアプリケーションを作成しました(Linux版を用意できなくて申し訳ないです)。アプリケーションが配布される形式は Topcoder のマラソンマッチ経験者なら慣れていそうですが、AHC はブラウザなのでそちらに慣れている方も多いかもしれません。Siv3D には web版も用意されているので、次回以降はそちらと併用することも検討しています。
グリッド上で敵機も自機もただの正方形で表示されているシンプルなグラフィックですが、それでも終盤に敵機をたった一撃で次々に破壊していく様子は見ていて爽快だったのではないでしょうか。
やはりこうしたアプリケーションを一番使い慣れている言語で実装できるのは助かりますね。この場を借りて Siv3D の制作関係者に御礼を申し上げます。

問題作成の経緯

ヒューリスティックの問題を作成する時はストーリーから作ることが多く、今回もそうでした。題材はいくつか考えていたのですが、シューティングゲームは問題を作りやすいかなと思って採用しました。最初から全ての敵機の情報がわかっている形式でも良いかなとも思っていましたが*2、より実際のゲームに近いリアクティブ形式にしてみました。敵機を倒し続けると自機がレベルアップして、ゲームが進むにつれて敵機も強くなっていく、というよくありそうなゲーム設定にしました。
ここまでは特に問題はなかったのですが、実際に問題を作るとなるとパラメータの設定がかなり難しいことに気付きました。敵機が弱すぎて簡単に殲滅できても困るし、逆に強すぎて全然倒せなかったらどうしようもありません。敵機の出現頻度や、自機と敵機の成長のバランスも悩ましかったです。とりあえず適当に設定して自分で解いていたのですが、後にビジュアライザが完成して結果を眺めてみるとかなりめちゃくちゃなパラメータ設定になっていました。やっぱりビジュアライザは偉大です。
色々とパラメータをいじってみたのですが、最終的には「序盤はパワーを重視して自機を強化、終盤で敵機を大量に破壊して一気にスコアを稼ぐ」といったゲームの性質の変化を実現できたかなと思います。

雑記

フィールドの左右がつながっていた？！

今回の問題はフィールドの左端と右端がつながっているという設定でした。最初はこの設定はなかったのですが、ずっと真ん中付近に陣取るのが強くなってしまいそうだなあと思って変更してみました。ここの説明文は特に強調していなかったこともあって、気付かなかった参加者が一定数いたようでした。ごめんなさい！

コンテストの形式について

yukicoder でスコアコンテストを開くのは3回目ですが、前回と同様にチームでの参加を可能としました。これはチームで参加できるヒューリスティックコンテストがかなり限られているので、そういった機会を提供したいという思いがありました*3。加えてコンテスト中でも問題への言及やビジュアライザの共有をほぼ制限しませんでした。普段の緊張感があるコンテストとは違いお祭り感覚で盛り上がってくれたら良いなと思っています。

優勝争い

順位表をずっと眺めていました。かなり早い時間帯から tomerun さんが高いスコアを出していてずっとリードしていたんですが、終盤にてりーさんが猛追してきて勝手に盛り上がっていました。ところが少し目を離した隙に nikaj さんという方がトップに躍り出ていて、なんか見たことないすごい人が現れたなーとか思っていました。よく考えると(よく考える前に気付くべきだった)、Topcoder のマラソンマッチに nika さんというレジェンドがいることに気付きます。いやまさかね、問題文は日本語しかないしさすがにこのコンテストの存在を知りようがないし、と思って恐る恐るコードを覗いてみたところ、どう見ても nika さんでした。もう yukicoder も国際化の時代です。
ここら辺で勝手に盛り上がりがピークに達していましたが、最後の提出で見事てりーさんが逆転優勝を決めました。世界トップレベルの戦い、とても見応えがありました。それにしてもてりーさん強すぎませんか？

最後に

今回も無事にコンテストを終えられてほっとしています。改めてテスターの yunix さん、コンテストに参加してくださった皆さん、参加できなかったけど問題を解いてくださった皆さん、ありがとうございました。そしてコンテストを開催する場を提供してくださっている yukicoder さん、いつもありがとうございます。これからもお世話になります。
今回はコンテスト終了後に twitter で振り返りスペースを開いてみました。そちらに参加してくださった皆さんもありがとうございました。
これからもコンテストを開いていきたいと思っていますので、次回以降もよろしくお願いします。

そういえば、スペースではマラソン作問が話題になりました。きっと次回の yukicoder score contest は yunix さんが開催してくださると信じて、今回は筆を置きたいと思います。

*1:ある程度行動を制限していたので完全にランダムにプレイアウトしたわけではないです。

*2:気付いた方もいらっしゃるかもしれませんが、問題のタイトルはいわゆるTASのもじりです。ツールを使って敵機の情報を抜き取ってスーパープレイをするという設定も考えたのですが、結局通常のプレイに落ち着きました。

*3:皆解会からも刺激を受けています。チームで色々と議論するのは勉強にもなりますし、やっぱり楽しいですね。

ヒューリスティックコンテストでベイス推定に入門しよう

ヒューリスティックコンテストでは、一部の数値が与えられないタイプの問題が出題されることがあります。多くの場合はインタラクティブ形式の問題で、少しずつ与えられる情報で数値を予測していくことになります。
(例：AHC003, HTTF2022予選, HTTF2023本選, AHC018)

こうした推定タイプの問題では、ベイズ推定が強力な手法となることがあります。
本記事はベイズ推定を使って問題を解いてみたい方の第一歩となることを目的としています(C++での実装例も載せました)。題材はAHC003としました。問題の概要は以下です。

  • 30×30のグリッド(上下左右が辺で結ばれている)があり、各辺の長さは与えられない。
  • 始点と終点が与えられるので経路を出力すると、その経路の長さ(±10%の誤差を含む)が返ってくる。これを1000回繰り返す。
  • 出力した経路が最短経路に近いほど得点が高くなる。回数を重ねるごとに得点の比重が上がる(つまり後半が大事)。
  • (入力で与えられない)真の辺の長さは以下の二つのパターンがある。
    • 縦横60本の同一直線上の辺は大体同じ長さ ($M=1$)
    • 各直線がランダムな位置で二分割されており、それぞれの側の辺は大体同じ長さ ($M=2$)

ということで、辺の長さを正確に推測できればダイクストラ法によって良い経路を導けそうです。

この記事は優勝者のeivourさんによる解説記事を参考にしています。
qiita.com
今回は行列演算の高速化やMCMCといった高度な内容は扱いませんので、線形代数と統計学の基礎知識があれば読めるのではないかと思います。より高い点数を目指したい方はeivourさんの記事をぜひご覧ください。
(筆者も勉強したばかりの内容も含んでいますので、厳密さに欠ける部分もあるかもしれません。不備がありましたらご指摘いただけますと幸いです。)

実際に問題に取り組む前にベイズ推定について説明します。既に理解している方は飛ばしてください。

ベイズ推定

ベイズ推定の基本となるベイズの定理について調べると、

$$P(A|B) = \frac{P(B|A)P(A)}{P(B)}$$

という数式を見ると思います。この式の意味は既に多くの解説があるのでここでは詳しく説明せず、早速具体例を見ていきましょう。
いきなりAHC003に挑戦しても良いのですが、その前に一つ簡単な例で感覚をつかむことにします。

  • 表が出る確率が未知のコインがある。このコインを $3$ 回投げたところ全て表だった。表が出る確率を推定しなさい。

普通は表が出た回数の割合としたいですが、そうすると表が出る確率 $p=1$ となってしまい、さすがにおかしい気がしませんか？
でもなぜおかしいと感じるのでしょうか。それはおそらく、現実のコインは普通は表と裏が出る確率にそれほど大きな差はないという前提知識があり、たまたま偏っただけだと考えるからだと思います。この感覚を数学的に表現できるのがベイズ推定です。
表が出る確率 $p$ をピンポイントで推定しても良いのですが、確率分布とすることもできます。そうすると推定の自信の度合いを分散という数学的な値で表現できて便利です。

では、表が出る確率(分布)を $P(A)$ としましょう。これが $B$ というコイントスの結果を反映し、条件付き確率 $P(A|B)$ に変化したと考えます。このように観測結果をもとに確率(分布)を更新する、というのがベイズ推定の基本的な考え方です。$P(A)$ は観測前の確率分布なので「事前分布」、$P(A|B)$ は観測後なので「事後分布」と呼びます。
ベイズの定理の式を眺めると、事後分布 $P(A|B)$ を求めるには $P(B|A)$ と $P(B)$ がわかれば良さそうです。ただ $P(B)$ は $A$ とは無関係(すなわちこのコインとは無関係)な定数*1なので後回しにして $P(B|A)$ を考えます。これは $A$ という条件のもとでの $B$ の確率、すなわち「表が出る確率が $p$ とわかっている場合に $B$ というコイントスの結果が得られる確率」のことですから、

$$P(B|A) = {}_3 \mathrm{C}_3\ p^3 (1-p)^0$$

です。一般にコインを $n$ 回投げて $k$ 回表が出たのならば、

$$P(B|A) = {}_n \mathrm{C}_k \ p^k (1-p)^{n-k}$$

となります。
今まで $P(A)$ の具体的な形は考えていませんでしたが、$P(B|A)$ と同じ形、すなわち

$$P(A) = Cp^{\alpha} (1-p)^{\beta}$$

のような形($C$ は定数で、この形の確率分布をベータ分布と呼びます)をしていれば、

$$P(A)P(B|A) = C'p^{\alpha + k} (1-p)^{\beta + n - k}$$

といったように $P(B|A)$ をかけた後も同じ形になりますね。事後分布が事前分布と同じ形の確率分布になるような事前分布のことを特に「共役事前分布」と呼びます。計算がとても楽になるので、可能なら共役事前分布を選びたいです。
なぜ計算が楽なのかと言うと、定数 $C'$ が簡単に求まるからです。事後分布もベータ分布だとわかっているなら、途中の定数の計算は省略して最後に求めれば良いです*2。つまり、共役事前分布ならば後回しにしていた $P(B)$ の計算は実は不要でした。

これで事後分布が求められるようになりました。最後に事前分布における $\alpha, \beta$ の値を具体的に設定してみましょう。実はこれは特に決まりはなく、設定する人の「感覚」が反映されます。「普通のコインなら、平均は $0.5$ で多少のばらつきがありそう」みたいな感覚です。例えば $\alpha = 50, \beta = 50$ とすると、平均が $0.5$、標準偏差が約 $0.050$ のベータ分布になります。
コインを3回投げて全て表だったという結果を反映すれば、$\alpha = 53, \beta = 50$ と更新されます。これは平均が約 $0.515$、標準偏差が約 $0.049$ のベータ分布です。たった3回の試行だったので平均はそれほど大きく動きませんでした(このコインは少し表が出やすいだろう程度でしょうか)。試しにコイントスの結果($n, k$ の値)や、事前分布の設定($\alpha, \beta$ の値)を色々変えて結果を確認してみてください。

AHC003に挑戦

それでは実際の問題に取り組んでみましょう。この記事では簡単のため $M=1$ とします(すなわち同一直線上の辺は全て大体同じ長さとします)。上のベイズ推定の項目での議論から、

  • 推定したい値の確率分布を設定する
  • 観測結果をもとに確率分布を更新する

という二つの作業が必要だとわかると思います。逆に言えばこの二つが適切にできれば高い性能を期待できます。

確率分布の設定

まず最初の確率分布を設定しましょう。どんな確率分布を選ぶべきかは問題設定によりますし、センスが問われるところでもあって難しいのですが、今回は正規分布を用いることにします*3。というのは、返ってくる経路の長さの誤差も正規分布に従うとしてしまえば、正規分布は共役事前分布になる(事前分布が正規分布なら事後分布も正規分布になる)から嬉しいわけです。
推定したい辺は全部で $1740$ 本ありますから、多変量正規分布になります。つまりそれぞれの辺の分散に加えて共分散も必要です。ただ変数が $1740$ 個もあると大変なので、とりあえずそれぞれの直線だけを考えて $60$ 変数としてしまいましょう。こうするとそれぞれの変数は独立になりますから、共分散は全て $0$ です。
辺の長さ $X$ は平均(ベクトル) $\boldsymbol{\mu}$、分散共分散行列 $V$ の正規分布に従うとします。各辺の平均は $5000$ としておくと分散 $v$ は、

$$v = \frac{1}{8000-2D} \int_{1000+D}^{9000-D} (x-5000)^2dx = \frac{1}{3} (4000-D)^2$$

と計算できます。$D$ の値がわからないですが、例えば {$250, 750, 1250, 1750$} の $4$ 通りを設定して全て計算しておき、後でどれが近そうかを推測することにしましょう。これで $\mu, V$ を具体的に設定できました(今回は共分散が全て $0$ ですから $V$ は対角行列です)。

事後分布の導出

次に、得られた観測結果から事後分布を求める方法を考えます。出力した経路においてその辺を通る回数*4を要素としたベクトル $\boldsymbol{c}$ を考えると、経路の長さの期待値 は $c^T \mu$ とベクトルの内積で表せます($c$ や $\mu$ は列ベクトルで、転置の記号 $T$ を用いて行ベクトルとしています)。そして、返ってきた経路の長さ $y$ は誤差項 $\varepsilon$ が加わった

$$y = c^T X + \varepsilon$$

という確率変数であるとします。今回の誤差は±10%でしたから、誤差項 $\varepsilon$ は $X$ とは独立で、平均 $0$、分散 $R=(0.1y)^2$ (標準偏差 $0.1y$) の正規分布に従うとしてみます。
この時の事後分布の導出は、こちらの記事(分散のみですがWoodburyの公式*5についても言及があります)やこちらの記事(一般の場合について導出を追えます)が大変参考になると思います。
ここでは結論だけを書きます。$y$ の分散を $v_y = c^T Vc + R$ とおくと、事後分布の平均 $\mu '$ および分散共分散行列 $V'$ は、

$$\mu ' = \mu + Vc\ v_{y}^{-1} (y-c^T \mu)$$
$$V' = V - Vc\ v_{y}^{-1}c^T V$$

となります。すなわち $X$ の事後分布は平均 $\mu '$、分散共分散行列 $V'$ の正規分布です。

$D$ の値はどれが適切なのか？

ここまでは勝手に $D$ の値を $4$ 通り設定して全てについて計算を進めていましたが、その計算結果からどの値が適切かを推測するにはどうしたら良いでしょうか。
$X$ を既知とした時に実際の観測結果 $y$ が得られる確率を比較すれば良さそうです(その確率があまりにも低いなら、そもそも仮定した $D$ の値が不適切だったと言えそうです)。この確率とは何かと言うと、「パラメータがわかっている時にその観測結果が得られる条件付き確率」ですから、上のベイズ推定の議論で言う $P(B|A)$ に他なりません。これを統計学の用語で「尤度」と呼びます。
確率変数 $y$ は平均 $c^T \mu$、分散 $v_y$ の正規分布に従うので、尤度は

$$p(y|X) = \frac{1}{\sqrt{2\pi v_y}} \exp(-\frac{(y-c^T \mu)^2}{2v_y})$$

です(意味合いは異なりますが、形は確率密度関数と同じです)。全ての観測結果の同時確率について尤度を考えたいので、これまでの尤度を全てかけ合わせます。しかしかけ算を繰り返すと値がほとんど $0$ になってしまって扱いにくい(例えば $e^{-100}$ とかになったりします)ので、実装する際には自然対数を取って足し合わせることにしましょう。そして、全ての $D$ の値について尤度を計算してその比で「もっともらしさ」を判断することにします。

経路の選択

ここまでで、$4$ 種類のモデルについて尤度と事後分布を求められるようになりました。あとはダイクストラ法で経路を決定したいのですが、そのためには各辺の長さを固定する必要があります。
モデルごとに考えます。平均ベクトル $\mu$ をそのまま長さとしても良いかもしれませんが、せっかく分散も計算してあるのでそれも活用したいです。分散が大きい場合はその辺についての推定が弱いということですから、積極的に経路に組み込んで情報を得たくなります。そこで、分散が大きいほど辺の長さが小さくなるようにしてみたくなります。例えば、平均から標準偏差の定数倍を引いた

$$w = \mu - \gamma \sqrt{v_x}$$

などが考えられそうです(分散 $v_x$ は $V$ の対角成分です)。分散の影響をどの程度反映するかは調整が可能なので、実際に色々な計算式を試してみてください。モデルごとに辺の長さを確定させたら、尤度の比で重みをつけて最終的な辺の長さとします。

C++での実装例

ここで示したのはあくまで例であり、機能としても最低限のものです。ここまでの議論が理解できた方はぜひ自力で実装してみてください。
前提として行列演算のライブラリが必要です。C++の場合はEigenという便利なライブラリがありますが、コンテストで使うために自分で実装してみます。

struct VectorXd {
    std::vector<double> vec;
    bool transposed = false;

    VectorXd() {}
    VectorXd(size_t __new_size) {
        vec.resize(__new_size);
    }
    VectorXd(size_t __new_size, double val) {
        vec.resize(__new_size, val);
    }
    void resize(size_t __new_size) {
        vec.resize(__new_size);
    }
    void resize(size_t __new_size, double val) {
        vec.resize(__new_size, val);
    }
    size_t size() const {
        return vec.size();
    }
    void emplace_back(const double val) {
        vec.emplace_back(val);
    }

    double& operator [] (const int num) {
        assert(num >= 0 && num < (int)vec.size());
        return vec[num];
    }
    double operator [] (const int num) const {
        assert(num >= 0 && num < (int)vec.size());
        return vec[num];        
    }
    VectorXd& operator += (const VectorXd& A) {
        assert(vec.size() == A.vec.size());
        for(int i = 0; i < (int)vec.size(); i++) {
            vec[i] += A.vec[i];
        }
        return *this;
    }
    VectorXd& operator -= (const VectorXd& A) {
        assert(vec.size() == A.vec.size());
        for(int i = 0; i < (int)vec.size(); i++) {
            vec[i] -= A.vec[i];
        }
        return *this;
    }
    VectorXd& operator *= (const double a) {
        for(auto& val : vec) {
            val *= a;
        }
        return *this;
    }
    VectorXd operator + (const VectorXd& A) const {
        assert(vec.size() == A.vec.size());
        VectorXd B = *this;
        return B += A;
    }
    VectorXd operator - (const VectorXd& A) const {
        assert(vec.size() == A.vec.size());
        VectorXd B = *this;
        return B -= A;
    }
    VectorXd operator * (const double a) const {
        VectorXd B = *this;
        return B *= a;
    }

    double dot (const VectorXd& A) const {
        assert(vec.size() == A.vec.size());
        double res = 0.0;
        for(int i = 0; i < (int)vec.size(); i++) {
            res += vec[i] * A.vec[i];
        }
        return res;
    }
    double squaredNorm() {
        return this->dot(*this);      
    }
    double norm() {
        return sqrt(this->squaredNorm());
    }

    void transpose() {
        transposed = !transposed;
    }
};
std::ostream& operator << (std::ostream& os, const VectorXd& A) {
    for(auto& val : A.vec) {
        os << val << " ";
    }
    return os;
}

struct MatrixXd {
    int row_size, column_size;
    std::vector<VectorXd> mat;
    std::vector<VectorXd> L, U;

    MatrixXd() {}
    MatrixXd(int r, int c) : row_size(r), column_size(c) {
        mat = std::vector(r, VectorXd(c));
    }
    MatrixXd(int r, int c, double val) : row_size(r), column_size(c) {
        mat = std::vector(r, VectorXd(c, val));
    }
    void resize(int r, int c) {
        row_size = r; column_size = c;
        mat.resize(r, VectorXd(c));
    }
    void resize(int r, int c, double val) {
        row_size = r; column_size = c;
        mat.resize(r, VectorXd(c, val));
    }

    int get_row_size() const {
        return row_size;
    }
    int get_column_size() const {
        return column_size;
    }

    VectorXd& operator [] (const int num) {
        assert(num >= 0 && num < (int)mat.size());
        return mat[num];
    }
    MatrixXd& operator += (const MatrixXd& A) {
        assert(row_size == A.row_size && column_size == A.column_size);
        for(int i = 0; i < row_size; i++) {
            for(int j = 0; j < column_size; j++) {
                mat[i][j] += A.mat[i][j];
            }
        }
        return *this;
    }
    MatrixXd& operator -= (const MatrixXd& A) {
        assert(row_size == A.row_size && column_size == A.column_size);
        for(int i = 0; i < row_size; i++) {
            for(int j = 0; j < column_size; j++) {
                mat[i][j] -= A.mat[i][j];
            }
        }
        return *this;
    }
    MatrixXd& operator *= (const MatrixXd& A) {
        assert(column_size == A.row_size);
        MatrixXd B(row_size, A.column_size);
        for(int i = 0; i < row_size; i++) {
            for(int k = 0; k < column_size; k++) {
                for(int j = 0; j < A.column_size; j++) {
                    B.mat[i][j] += mat[i][k] * A.mat[k][j];
                }
            }
        }
        return *this = B;
    }
    MatrixXd& operator *= (const double a) {
        for(int i = 0; i < row_size; i++) {
            for(int j = 0; j < column_size; j++) {
                mat[i][j] *= a;
            }
        }
        return *this;
    }
    MatrixXd operator + (const MatrixXd& A) const {
        assert(row_size == A.row_size && column_size == A.column_size);
        MatrixXd B = *this;
        return B += A;
    }
    MatrixXd operator - (const MatrixXd& A) const {
        assert(row_size == A.row_size && column_size == A.column_size);
        MatrixXd B = *this;
        return B -= A;
    }
    MatrixXd operator * (const MatrixXd& A) const {
        assert(column_size == A.row_size);
        MatrixXd B = *this;
        return B *= A;
    }
    MatrixXd operator * (const double k) const {
        MatrixXd B = *this;
        return B *= k;
    }
    VectorXd operator * (const VectorXd& A) const {
        assert(column_size == (int)A.size());
        VectorXd B(row_size);
        for(int i = 0; i < row_size; i++) {
            for(int j = 0; j < column_size; j++) {
                B[i] += mat[i][j] * A[j];
            }
        }
        return B;
    }
    MatrixXd pow(const long long n) const {
        assert(row_size == column_size);
        if(n == 0){
            MatrixXd E(row_size, column_size);
            for(int i = 0; i < row_size; i++) {
                E.mat[i][i] = 1.0;
            }
            return E;
        }
        MatrixXd A = pow(n >> 1);
        A *= A;
        if(n & 1) A *= *this;
        return A;
    }

    VectorXd row(const int num) const {
        assert(num >= 0 && num < row_size);
        return mat[num];
    }
    VectorXd col(const int num) const {
        assert(num >= 0 && num < column_size);
        VectorXd A(row_size);
        for(int i = 0; i < row_size; i++) {
            A[i] = mat[i][num];
        }
        return A;
    }
    void row_apply(const int num, const VectorXd& A) {
        assert(num >= 0 && num < row_size);
        mat[num] = A;
    }
    void col_apply(const int num, const VectorXd& A) {
        assert(num >= 0 && num < column_size);
        for(int i = 0; i < row_size; i++) {
            mat[i][num] = A[i];
        }
    }
    void row_append(const VectorXd& A) {
        assert(column_size == (int)A.size());
        row_size++;
        mat.emplace_back(A);
    }
    void column_append(const VectorXd& A) {
        assert(row_size == (int)A.size());
        column_size++;
        for(int i = 0; i < row_size; i++) {
            mat[i].emplace_back(A[i]);
        }
    }
    VectorXd diag() {
        assert(row_size == column_size);
        VectorXd A(row_size);
        for(int i = 0; i < row_size; i++) {
            A[i] = mat[i][i];
        }
        return A;
    }
    MatrixXd transpose() {
        MatrixXd A(column_size, row_size);
        for(int i = 0; i < column_size; i++) {
            for(int j = 0; j < row_size; j++) {
                A[i][j] = mat[j][i];
            }
        }
        return A;
    }

    void LU_decomposition() {
        assert(row_size == column_size);
        int N = row_size;
        U = mat;
        L.resize(N, VectorXd(N));
        for(int i = 0; i < N; i++) {
            L[i][i] = 1.0;
        }
        for(int i = 0; i < N; i++) {
            for(int j = i + 1; j < N; j++) {
                double val = U[j][i] / U[i][i];
                for(int k = i; k < N; k++) {
                    U[j][k] -= U[i][k] * val;
                }
                L[j][i] += val;
            }
        }
    }
    MatrixXd inverse() {
        assert(row_size == column_size);
        int N = row_size;
        MatrixXd A = *this;
        MatrixXd E(N, N);
        for(int i = 0; i < N; i++) {
            E[i][i] = 1.0;
        }
        for(int i = 0; i < N; i++) {
            double inv = 1.0 / A[i][i];
            for(int j = 0; j < N; j++) {
                A[i][j] *= inv;
                E[i][j] *= inv;
            }
            for(int j = 0; j < N; j++) {
                if(i == j) continue;
                double val = A[j][i];
                for(int k = 0; k < N; k++) {
                    A[j][k] -= A[i][k] * val;
                    E[j][k] -= E[i][k] * val;
                }
            }
        }
        return E;
    }
    VectorXd inverse_product(const VectorXd& A) const {
        assert(row_size == column_size);
        int N = row_size;
        assert((int)A.size() == N);
        VectorXd Y(N);
        for(int i = 0; i < N; i++) {
            double sum = 0.0;
            for(int j = 0; j < i; j++) {
                sum += L[i][j] * Y[j];
            }
            double res = (A[i] - sum) / L[i][i];
            Y[i] = res;
        }
        VectorXd X(N);
        for(int i = N - 1; i >= 0; i--) {
            double sum = 0.0;
            for(int j = N - 1; j > i; j--) {
                sum += U[i][j] * X[j];
            }
            double res = (Y[i] - sum) / U[i][i];
            X[i] = res;
        }
        return X;
    }
};
std::ostream& operator << (std::ostream& os, const MatrixXd& A) {
    for(int i = 0; i < A.row_size; i++) {
        os << A.row(i) << std::endl;
    }
    return os;
}

もっと便利にできるとは思いますが、ひとまずこれで挑みます。(LU分解や逆行列の計算は今回は使いません)
準備ができたので、ベイズ推定のクラスを作ってみます。(これもかなり洗練の余地が残っていると思います)

struct distribution {
    double mu, var;
    distribution(double mu, double var) : mu(mu), var(var) {}
};
struct multi_distribution {
    VectorXd mu;
    MatrixXd cov;
    multi_distribution(VectorXd& mu, MatrixXd& cov) {
        this->mu = mu;
        this->cov = cov;
    }
};
class Bayesian {
private:
    multi_distribution md;
    double log_likelihood;
    double l_sum;

public:
    Bayesian(multi_distribution& md) : md(md) {
        log_likelihood = 0.0;
        l_sum = 0.0;
    }

    void update(const VectorXd& c, double y, const distribution& err) {
        double Syy = c.dot(md.cov * c) + err.var;
        VectorXd Sxy = md.cov * c;
        int siz = c.size();
        MatrixXd res(siz, siz);
        for(int i = 0; i < siz; i++) {
            for(int j = 0; j < siz; j++) {
                res[i][j] = Sxy[i] * Sxy[j];
            }
        }
        res *= (1.0 / Syy);
        double dif = y - c.dot(md.mu);

        md.cov -= res;
        md.mu += Sxy * dif * (1.0 / Syy);
        log_likelihood = -0.5 * (log(2.0 * M_PI * Syy) + dif * dif / Syy);
        l_sum += log_likelihood;
    }

    VectorXd get_mu() const {
        return md.mu;
    }
    MatrixXd get_cov() const {
        return md.cov;
    }
    double get_likelihood() const {
        return log_likelihood;
    }
    double get_sum_likelihood() const {
        return l_sum;
    }
};

平均と分散共分散行列を指定した多変量分布でオブジェクトを作り、update 関数で確率分布や対数尤度を更新します。続いて、以下の calc_cost 関数で各モデルの事後分布から各辺の重みを計算します(INFは適当な巨大な数としてください)。辺の番号は、横の辺を上から順に $0-29$、縦の辺を左から順に $30-59$ としています。

VectorXd calc_cost(const std::vector<Bayesian>& bayes, int turn) {
    VectorXd w(29 * 60);
    int siz = bayes.size();
    double tot_l = 0.0;
    double max_likelihood = -INF;
    for(int m = 0; m < siz; m++) {
        if(bayes[m].get_sum_likelihood() > max_likelihood) {
            max_likelihood = bayes[m].get_sum_likelihood();
        }
    }
    for(int m = 0; m < siz; m++) {
        tot_l += exp(bayes[m].get_sum_likelihood() - max_likelihood);
    }
    for(int m = 0; m < siz; m++) {
        VectorXd cost = bayes[m].get_mu();
        VectorXd dev = bayes[m].get_cov().diag();
        for(int i = 0; i < (int)dev.size(); i++) {
            dev[i] = sqrt(dev[i]);
        }
        double weight = 1.0 - (double)turn / 1000.0;
        cost -= dev * weight;
        cost *= exp(bayes[m].get_sum_likelihood() - max_likelihood) / tot_l;
        for(int i = 0; i < 29 * 60; i++) {
            w[i] += cost[i/29];
        }
    }
    for(int i = 0; i < 29 * 60; i++) {
        w[i] = std::max(w[i], 1000.0);
        w[i] = std::min(w[i], 9000.0);
    }
    return w;
}

辺の長さの計算にあたり、標準偏差にかかる定数を $\gamma = 1.0 - t/1000$ としてみました。ここで $t$ は試行回数です。また、長さを小さく見積もりすぎて負の値になってしまうとダイクストラ法が使えなくなるので、$1000$ 以上 $9000$ 以下という制約を設けることにしました。
尤度の比を計算する際に注意点があります。対数尤度から尤度に直すには指数関数を用いれば良いですが、前述の通り尤度はほとんど $0$ であるため計算できない可能性があります。重要なのは比だけですから、あらかじめそれぞれの対数尤度に適当な値を足してから比を計算するようにしています。これで辺の長さが求まったので、ダイクストラ法で経路を決めることができます。
以上の実装で提出したところ、約94.4%の得点が得られました。コンテストの順位では上位10%より少し下くらいでしょうか。しかし変数を $60$ 個に減らしたり $M=1$ のみを考えたりと制限が多いですから、もっと高い点数を目指せそうです。

どんな場合にベイズ推定が有効なのか？

推定タイプの問題と言っても様々なパターンがあります。ベイズ推定を活用するには、繰り返しにはなりますが

  • 推定したい値の確率分布を適切に設定できる
  • 観測結果をもとに確率分布を更新できる

ことが重要です。正規分布が便利なことが多いですが、平均と分散共分散行列を適切に決められる必要があります。また、今回は $y = c^T X + \varepsilon$ という線形の関係式を仮定できましたが、関係式がよくわからない場合はどうするのか、という問題もありそうです。

ベイズ推定を応用したガウス過程回帰という手法があります。誤解を恐れずにすごく簡単に説明すると、新しい入力が与えられた時、今までの入力との"距離(どれくらい似ているか)"で重み付けをして出力を予測する手法です。関係式がわからなくても、新しい入力と今までの入力との"距離"を適切に設定できそうなら、ガウス過程回帰を活用できるかもしれません。実際にAHC018ではガウス過程回帰を利用した解法がありました。興味がある方はてりーさんの解説記事やyosssさんの解説スライドをぜひご覧ください。ガウス過程回帰についてはここではこれ以上触れませんが、いつか記事を書くことがあるかもしれません。

長くなってしまいましたが、最後までお読みいただきありがとうございました。数字をいじって結果を眺めて遊んでみたり、色々な問題でベイズ推定を活用できないか考えてみたりして、ベイズ推定と仲良くなりましょう。

*1:これはコインによらず一般にこの結果が得られる確率のことなので、表が出る確率を $p$ と置いて $p$ の全区間で積分して求めることになります。この積分は計算が困難な場合もあるので、計算しなくて済むと大変嬉しいです。

*2:確率分布は全区間で積分して $1$ になる必要があるので、事後分布の積分さえ計算できれば良いです。有名な確率分布の場合は積分の結果が既に知られていることが多いです。その上、そもそも定数を求めなくても確率分布の平均と分散は計算できてしまいます。

*3:問題の入力生成の項目を見ると、辺の長さは本来正規分布に従うわけではなさそうです。とはいえ正規分布を仮定しておけば、試行回数を重ねるごとに、平均は真の辺の長さに近づき分散は小さくなった正規分布が得られると期待できますから、悪くない仮定のように思えます。ただし $M=2$ の場合は簡単ではないです。詳しくはeivourさんの記事をご覧ください。

*4:今回は $60$ 変数としたため同一直線上の辺を全て同一視していますから、同じ辺を複数回通ることがあり得ます。全ての辺を考えた $1740$ 変数のモデルならば、ベクトル $c$ の各要素は $1$ または $0$ をとります。

*5:導出方法によっては $V' = (V^{-1} + cR^{-1} c^T)^{-1}$ となっているかもしれませんが、Woodburyの公式より記事中の式と等しいです。Woodburyの公式についてやどちらの式が望ましいかは記事中のリンクをご覧ください。

AtCoder Beginner Contest 158 F - Removing Robots

ABC158-F問題です。公式解説とだいぶ違う方針だったので書き残したいと思います。

問題概要：

数直線上に$\ N\ $個のロボットがあり、$i\ $番目のロボットは座標$\ X_i\ $にいる。各ロボットは起動すると右に$\ D_i\ $だけ動く。他のロボットにぶつかるとそのロボットも起動する。移動を終えたロボットは取り除かれる。
高橋君は好きな回数ロボットを起動できる。数直線上に残るロボットの組み合わせは何通り考えられるか？
・$1 \leq N \leq 2 \times 10^5$
・$-10^9 \leq X_i \leq 10^9$
・$1 \leq D_i \leq 10^9$

ロボットを起動する順番は重要ではなさそうです。そこで座標の昇順に次のようなDPを考えたくなります。

$dp[i]\ :=\ $左から$\ i\ $番目のロボットまで見たときの組み合わせの総数

各ロボットについて、「そのロボットを起動するかしないか」の$\ 2\ $通りがあるので、このDPは単に総数が倍になっていくだけになってしまいます。しかし実際には、前のロボットの動き次第ではそのロボットは強制的に起動するので、これだけでは不十分とわかります。
では一つ前のロボットが起動したかどうかも状態に持てば良いのではないか、という発想になります。しかしこれでもまだ不十分で、もっと前のロボットに巻き込まれる場合もあります。
重要なのは前のロボットが動いたかどうかではなく、「一つ前のロボットまでで、最も右に来たロボットがどこまで来たのか」です。そう考えると、次のようなDPが浮かびます。

$dp[i][j]\ := \ $左から$\ i\ $番目のロボットまで見て、座標$\ j\ $までロボットが移動するような組み合わせの総数

座標は非常に大きな値を取りうるので、map(あるいはunordered_map)で管理することにします。
遷移がどうなるかを考えてみます。

・$j \leq X_{i+1}\ $の場合：ロボット$\ i+1\ $を起動するかしないか選べます。起動しない場合の右端は$\ X_{i+1},\ $する場合の右端は$\ X_{i+1} + D_{i+1}\ $です。
・$j > X_{i+1}\ $の場合：ロボット$\ i+1\ $は強制的に起動するので、右端は$\ max(j, X_{i+1} + D_{i+1})\ $です。

各ロボットについて$\ 3\ $通りの遷移しかなく、十分高速です。

行列構造体を使ってみる

競技プログラミングにおいて、行列演算が活躍する場面はそれなりにあります。
もちろんその都度配列を用意しても良いのですが特に行列の積の計算などは面倒で、やっぱり構造体で定義したくなりました。
二次元の行列で簡単な演算ができるように、構造体を書いてみたので以下に貼ります。
私は普段C++を使っていますがそこまで詳しくはないので、何かおかしな点に気付いた場合は教えてくださると幸いです。
(こちらの記事も参考になると思います)
競技プログラミング用にC++で行列（2次元平面）ライブラリを作成する - Qiita
作ったばかりなのでもしかしたらバグが潜んでいるかもしれないです。もし使用する場合は挙動の確認をお願いします。

ソースコード

template<typename T>
struct matrix{
	int row, column;
	vector<vector<T>> mat;

	matrix(int r = 0, int c = 0, T val = 0) : row(r), column(c) {
		mat = vector<vector<T>>(r, vector<T>(c, val));
	}
	vector<T>& operator [] (const int num) {
		return mat[num];
	}
	matrix& operator += (const matrix A) {
		assert(row == A.row && column == A.column);
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				mat[i][j] += A.mat[i][j];
			}
		}
		return *this;
	}
	matrix& operator -= (const matrix A) {
		assert(row == A.row && column == A.column);
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				mat[i][j] -= A.mat[i][j];
			}
		}
		return *this;
	}
	matrix& operator *= (const matrix A) {
		assert(column == A.row);
		matrix B(row, A.column);
		for(int i = 0; i < row; i++) {
			for(int k = 0; k < column; k++) {
				for(int j = 0; j < A.column; j++) {
					B.mat[i][j] += mat[i][k] * A.mat[k][j];
				}
			}
		}
		return *this = B;
	}
	matrix& operator *= (long long k) {
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				mat[i][j] *= k;
			}
		}
		return *this;
	}
	matrix operator + (const matrix A) const {
		assert(row == A.row && column == A.column);
		matrix B = *this;
		return B += A;
	}
	matrix operator - (const matrix A) const {
		assert(row == A.row && column == A.column);
		matrix B = *this;
		return B -= A;
	}
	matrix operator * (const matrix A) const {
		assert(column == A.row);
		matrix B = *this;
		return B *= A;
	}
	matrix operator * (long long k) const {
		matrix B = *this;
		return B *= k;
	}
	vector<T> operator * (const vector<int> A) const {
		int vector_size = A.size();
		assert(column == vector_size);
		vector<T> B(row);
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				B[i] += mat[i][j] * A[j];
			}
		}
		return B;
	}
	matrix pow(long long n) const {
		assert(row == column);
		if(n == 0){
			matrix E(row, column);
			for(int i = 0; i < row; i++) {
				E.mat[i][i] = 1;
			}
			return E;
		}
		matrix A = pow(n >> 1);
		A *= A;
		if(n & 1) A *= *this;
		return A;
	}
};

簡単に説明すると、行列同士の和、差、積の計算、行列の定数倍やベクトルとの積、累乗計算などを用意しました。
行列同士の演算では、行数や列数の関係で定義されない演算をしようとすると実行時エラーが発生するようになっているはずです。
使用例として、この構造体を使ってABC189E - Rotate and Flipを解いてみます。
解説にもあるように各変換はアフィン変換ですので、3 x 3 行列で表せます。こういう時に構造体があると楽そうですね。

ACコード

#include <bits/stdc++.h>
#define rep(i,n) for(int i = 0; i < (int)(n); i++)

using namespace std;
using LL = long long;

template<typename T>
struct matrix{
	int row, column;
	vector<vector<T>> mat;

	matrix(int r = 0, int c = 0, T val = 0) : row(r), column(c) {
		mat = vector<vector<T>>(r, vector<T>(c, val));
	}
	vector<T>& operator [] (const int num) {
		return mat[num];
	}
	matrix& operator += (const matrix A) {
		assert(row == A.row && column == A.column);
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				mat[i][j] += A.mat[i][j];
			}
		}
		return *this;
	}
	matrix& operator -= (const matrix A) {
		assert(row == A.row && column == A.column);
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				mat[i][j] -= A.mat[i][j];
			}
		}
		return *this;
	}
	matrix& operator *= (const matrix A) {
		assert(column == A.row);
		matrix B(row, A.column);
		for(int i = 0; i < row; i++) {
			for(int k = 0; k < column; k++) {
				for(int j = 0; j < A.column; j++) {
					B.mat[i][j] += mat[i][k] * A.mat[k][j];
				}
			}
		}
		return *this = B;
	}
	matrix& operator *= (long long k) {
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				mat[i][j] *= k;
			}
		}
		return *this;
	}
	matrix operator + (const matrix A) const {
		assert(row == A.row && column == A.column);
		matrix B = *this;
		return B += A;
	}
	matrix operator - (const matrix A) const {
		assert(row == A.row && column == A.column);
		matrix B = *this;
		return B -= A;
	}
	matrix operator * (const matrix A) const {
		assert(column == A.row);
		matrix B = *this;
		return B *= A;
	}
	matrix operator * (long long k) const {
		matrix B = *this;
		return B *= k;
	}
	vector<T> operator * (const vector<T> A) const {
		int vector_size = A.size();
		assert(column == vector_size);
		vector<T> B(row);
		for(int i = 0; i < row; i++) {
			for(int j = 0; j < column; j++) {
				B[i] += mat[i][j] * A[j];
			}
		}
		return B;
	}
	matrix pow(long long n) const {
		assert(row == column);
		if(n == 0){
			matrix E(row, column);
			for(int i = 0; i < row; i++) {
				E.mat[i][i] = 1;
			}
			return E;
		}
		matrix A = pow(n >> 1);
		A *= A;
		if(n & 1) A *= *this;
		return A;
	}
};

int main(){
	int N;
	cin >> N;
	vector<LL> X(N), Y(N);
	rep(i,N) cin >> X[i] >> Y[i];
	int M;
	cin >> M;
	vector<matrix<LL>> cum(M + 1);
	rep(i,M+1) cum[i] = matrix<LL>(3, 3);
	rep(i,3) cum[0][i][i] = 1;
	rep(i,M){
		matrix<LL> A(3, 3);
		int t;
		cin >> t;
		if(t == 1){
			A[0][1] = 1, A[1][0] = -1, A[2][2] = 1;
		}
		else if(t == 2){
			A[0][1] = -1, A[1][0] = 1, A[2][2] = 1;
		}
		else if(t == 3){
			int p;
			cin >> p;
			A[0][0] = -1, A[0][2] = 2 * p, A[1][1] = 1, A[2][2] = 1;
		}
		else{
			int p;
			cin >> p;
			A[0][0] = 1, A[1][1] = -1, A[1][2] = 2 * p, A[2][2] = 1;			
		}
		cum[i+1] = A * cum[i];
	}
	int Q;
	cin >> Q;
	rep(i,Q){
		int a, b;
		cin >> a >> b;
		b--;
		vector<LL> p(3);
		p[0] = X[b], p[1] = Y[b], p[2] = 1;
		vector<LL> ans = cum[a] * p;
		cout << ans[0] << " " << ans[1] << endl;
	}

	return 0;
}

AtCoder Beginner Contest 171 F - Strivore

ABC171のF問題からです。このアプローチ(特に後半部分)を見かけなかったので書いてみました。

問題概要：

文字列$\ S\ $の好きな位置に好きな英小文字を$\ K\ $回挿入して得られる文字列は何種類あるか？
・$|S| \leq 10^6$
・$1 \leq K \leq 10^6$

まずは小さい例で実験してみます。例えば、$S=ab, K=1\ $とします。
どこに挿入するかで$\ 3\ $通り、何を挿入するかで$\ 26\ $通りあるので全部で$\ 78\ $通りありそうです。ただ、このままだと重複がありそうなので、どのような文字列を複数回数えているか考えます。
挿入したい文字をその文字と同じ文字の隣に入れる場合、左に入れても右に入れてもできる文字列は同じです。例えば$\ S=ab\ $に新たに$\ a\ $を挿入する場合、一番左に入れても真ん中に入れてもできる文字列は$\ aab\ $です。これは二文字目の$\ b\ $についても同じことが言えるので、答えは$\ 78-2=76\ $通りです。
実は上の議論において、元の文字列の中身は関係ないです。$S=aa\ $のように同じ文字が連続する場合でも、$\ aaa\ $を$\ 3\ $回数えるので答えはやはり$\ 78-2=76\ $通りです。
つまり、与えられる文字列$\ S\ $の中身は実は関係なく、その長さだけが重要なのでした。以降は$\ |S|=N\ $とします。
長さだけを考えれば良いのですから、与えられた文字列は$\ a\ $が$\ N\ $個連続しているものとします。そうすると一気に視界が良くなります。
ここに文字を$\ K\ $個挿入します。挿入する文字に$\ a\ $が含まれていると厄介そうなので、まずは全てそれ以外の文字とします。そうすると、$N\ $個の$\ a\ $と$\ K\ $個のそれ以外の文字の並べ替えになるので、全部で$\ {}_{N+K} \mathrm{ C }_N \times 25^K\ $通りになります。
($\ {}_{N+K} \mathrm{ C }_N\ $の方は$\ a\ $の位置の決め方、$25^K\ $の方は$\ a\ $以外の位置に何を挿入するか、に対応します)
こう考えると、挿入する文字に$\ a\ $が含まれている場合も同様に処理できます。すなわち挿入する文字に$\ a\ $が$\ i\ $個含まれているとし、全て足せば良いです。求める答えは
$$
\sum_{i=0}^{K} {}_{N+K} \mathrm{ C }_{N\\+i} \times 25^{K-i}
$$となります。
この式はDPなどから変形しても得られるようですが、このような見方をすると自然に導出できるような気がしました。