modint構造体の工夫 - とこはるのまとめ

(注 (2021/1/31) : この記事はとても古く、なおかつもっといい実装が考えられます。さらに言えば、AtCoderではac-libraryが整備されている現在では価値は無いと思われます)

はじめに

プログラミングコンテストの問題では、答えをある値で割ったあまりを求める問題が頻出ですが、バグを生じやすい特徴があります。その例は次の記事の序文に詳しく記述されています。
noshi91.hatenablog.com

この記事では、C++においてこの種の問題を解決するための対策として、modint構造体を導入しました。これにより、本質的な式の記述に集中することができるようになるなど、多くの利点を享受することができます。
しかし、次の欠点があることにも触れられています

実行時にMODの値が決定する場合は使用不可。代わりに別の実装を用いる必要がある。
アルゴリズムに依存した高速化 : 例えば「+= (intの範囲の数)」を行いたいときに途中でMODを取ることは時間の無駄ですが、そのチューニングが難しくなる

今回の記事では、これらの欠点をそれなりに克服するmodint構造体を紹介します。

コンセプト

具体的な演算を除いたメイン部分のコードを示します。

vector<long long> MODS = {998244353};
template <int kind = 0, int fast = 0>
class mint {
 public:
  long long v;

  mint() : v(0) {}
  mint(long long v)
      : v(fast == 0 ? (v < 0 ? (v % MODS[kind] + MODS[kind]) % MODS[kind]
                             : (v >= MODS[kind] ? v % MODS[kind] : v))
                    : (v)) {}
  long long get_mod() { return MODS[kind]; }
  long long get_val() { return v; }
  void take_mod() { v %= MODS[kind]; }
};

実行時にMODを指定する

特徴の一つはMODの値を可変値にすることです。この指定方法は途中で値に変更があると計算が壊れてしまう可能性があるという意味で優れたコードだとは言えません。しかしながら、実行時にMODを指定できるようにしたいため、このような書き方になりました。(もっときれいな書き方があるとは思いますので、あれば是非教えてください。)

テンプレート引数kindの初期値は0ですので、mint<>の宣言でデフォルト値のMODでのmodint型を宣言できます。また、ほかのMOD値を使用したい場合にはvectorコンテナMODSへ要素を追加し、mint<1>などの宣言をすれば使用できます。

また、一般にテンプレートを使用するメリットとして、異なるテンプレート値を持つオブジェクトの間の演算をコントロールできる点が挙げられます。

例えばmodint構造体では、積のような演算は同じMOD値でのみ行いたいという自然な要請があるため、異なるMOD値を持つオブジェクトどうしの演算はコンパイルエラーで弾いてほしいものです。これを実現するために、積の演算は次のように記述しています。

template <int kind>
mint<kind>& operator*=(mint<kind>& a, mint<kind> b) {
  return a = a.v * b.v;
}

これで同じMOD値(より正確にはテンプレートパラメータkindの値)どうしの演算でなければコンパイルエラーが発生することになり、誤答を未然に防ぐことができます。

また、同じMOD値でも本質的に混ぜてはいけなければ、同じMOD値を用いつつテンプレートパラメータkindを変えてもよいかもしれません（が、そんなことが必要な状況に出会ったことはありません）

最後に、コード例として、Educational DP Contestの問題S - Digit Sumへの提出Submission #5077234 - Educational DP Contestを紹介します。

この問題では数え上げに対するMOD値とDPのindexに関するMOD値の二つのMOD値を扱います。また、indexに関するMODは実行時に与えられています。先ほど紹介したコードではこれらを共通のクラスで処理出来ていることがわかります。

高速化オプション

次に、高速化オプションについて説明します。冒頭のコードのテンプレートパラメータにfastというものがありました。これは0のときは通常モードで、1のときはMOD演算をスキップすることで多少の高速化を行うモードであることを表しています。(現状テンプレートパラメータの型がintになっていますが、boolにする方がよさそうですね)

このパラメータを導入する理由は、冒頭でも説明した通り、intの範囲に収まる加算のみを行う場合には、速度を考慮に入れるとMODを取る回数を少なくしたい場合があるからです。

実際の例を紹介します。天下一プログラマーコンテスト2019の問題
F - Banned Xへの、fastパラメータをOFFにしている提出
Submission #5078264 - Tenka1 Programmer Contest 2019(341ms)
とONにしている提出
Submission #5074374 - Tenka1 Programmer Contest 2019(265ms)
です。

さて、この二つのコードの相違は"mint<> ans = 0;"と"mint<0, 1> ans = 0;"の違いのみです。したがって、この方針のメリットとして、一度プログラムがTLEしてしまった場合に、大規模な修正をかけることなく、演算手法を指定することができる点が挙げられます*1。

今回できていないこと

他の方に教えていただきましたが、コンパイラに定数であることが伝わっていればMOD演算が高速になります。
そこをサポートできるとよりよさそうです。

まとめ

modint構造体に対する2点の工夫を紹介しました
- 実行時にMOD値を指定する
- 高速化オプションのパラメータを加える
書き方はまだ改善ができそうに感じています
より良い書き方があればぜひ教えてください

*1:なお、このコードではfast=1に対する演算は加算しかサポートしていません。減算にも対応しているとよいなぁ、と書いている途中に思いました