src/util/probability-table.h

   1 #pragma once
   2
   3 #include "term/z-rand.h"
   4
   5 #include <algorithm>
   6 #include <exception>
   7 #include <stdexcept>
   8 #include <tuple>
   9 #include <vector>
  10
  11 /**
  12  * @brief 確率テーブルクラス
  13  *
  14  * 確率テーブルを作成し、確率に従った抽選を行うクラス
  15  *
  16  * @tparam IdType 確率テーブルに登録するIDの型
  17  */
  18 template <typename IdType>
  19 class ProbabilityTable {
  20 public:
  21     /**
  22      * @brief コンストラクタ
  23      *
  24      * 空の確率テーブルを生成する
  25      */
  26     ProbabilityTable() = default;
  27
  28     /**
  29      * @brief 確率テーブルを空にする
  30      */
  31     void clear()
  32     {
  33         return item_list_.clear();
  34     }
  35
  36     /**
  37      * @brief 確率テーブルに項目を登録する
  38      *
  39      * 確率テーブルに項目のIDと選択確率のセットを登録する。
  40      * 追加した項目が選択される確率は、
  41      * 追加した項目の確率(引数prob) / すべての項目のprobの合計
  42      * となる。
  43      * probが0もしくは負数の場合はなにも登録しない。
  44      *
  45      * @param id 項目のID
  46      * @param prob 項目の選択確率
  47      */
  48     void entry_item(IdType id, int prob)
  49     {
  50         if (prob > 0) {
  51             // 二分探索を行うため、probは累積値を格納する
  52             auto cumulative_prob = item_list_.empty() ? 0 : std::get<1>(item_list_.back());
  53             item_list_.emplace_back(id, cumulative_prob + prob);
  54         }
  55     }
  56
  57     /**
  58      * @brief 現在の確率テーブルのすべての項目の選択確率の合計を取得する
  59      *
  60      * 確率テーブルに登録されているすべての項目の確率(登録時の引数prob)の合計を取得する。
  61      *
  62      * @return int 現在の確率テーブルのすべての項目の選択確率の合計
  63      */
  64     int total_prob() const
  65     {
  66         if (item_list_.empty())
  67             return 0;
  68
  69         return std::get<1>(item_list_.back());
  70     }
  71
  72     /**
  73      * @brief 確率テーブルに登録されている項目の数を取得する
  74      *
  75      * 確率テーブルに登録されている項目の数を取得する。
  76      * 登録されている項目のIDに被りがある場合は別の項目として加算した数となる。
  77      *
  78      * @return size_t 確率テーブルに登録されている項目の数
  79      */
  80     size_t item_count() const
  81     {
  82         return item_list_.size();
  83     }
  84
  85     /**
  86      * @brief 確率テーブルの項目が空かどうかを調べる
  87      *
  88      * @return bool 確率テーブルに項目が一つも登録されておらず空であれば true
  89      *              項目が一つ以上登録されており空でなければ false
  90      */
  91     bool empty() const
  92     {
  93         return item_list_.empty();
  94     }
  95
  96     /**
  97      * @brief 確率テーブルから項目をランダムに1つ選択する
  98      *
  99      * 確率テーブルに登録されているすべての項目から、確率に従って項目を1つ選択する
 100      * それぞれの項目が選択される確率は、確率設定probに対して
 101      * 該当項目のprob / すべての項目のprobの合計
 102      * となる。
 103      * 抽選は独立試行で行われ、選択された項目がテーブルから取り除かれる事はない。
 104      * 確率テーブルになにも登録されていない場合、std::runtime_error例外を送出する。
 105      *
 106      * @return int 選択された項目のID
 107      */
 108     IdType pick_one_at_random() const
 109     {
 110         if (empty()) {
 111             throw std::runtime_error("There is no entry in the probability table.");
 112         }
 113
 114         // probの合計の範囲からランダムでkeyを取得し、二分探索で選択する項目を決定する
 115         const int key = randint0(total_prob());
 116         auto it = std::partition_point(item_list_.begin(), item_list_.end(), [key](const auto &i) { return std::get<1>(i) <= key; });
 117
 118         return std::get<0>(*it);
 119     }
 120
 121     /**
 122      * @brief 確率テーブルから複数回抽選する
 123      *
 124      * 確率テーブルから引数 n で指定した回数抽選し、抽選の結果選択された項目のIDを
 125      * 出力イテレータに n 個書き込む。
 126      * 抽選は独立試行で行われ、選択された項目がテーブルから取り除かれる事はない。
 127      *
 128      * @tparam OutputIter 出力イテレータの型
 129      * @param first 結果を書き込む出力イテレータ
 130      * @param table 抽選を行う確率テーブル
 131      * @param n 抽選を行う回数
 132      */
 133     template <typename OutputIter>
 134     static void lottery(OutputIter first, const ProbabilityTable &table, size_t n)
 135     {
 136         std::generate_n(first, n, [&table] { return table.pick_one_at_random(); });
 137     }
 138
 139 private:
 140     /** 項目のIDと確率のセットを格納する配列 */
 141     std::vector<std::tuple<IdType, int>> item_list_;
 142 };