モノイドの積を求めるやつ
セグ木とか Disjoint Sparse Table の話ではないです。
正式名称がわからないアルゴリズムも好き好き大好き（リサーチ不足）
導入
floor sum のモノイド積版です。丸紅プログラミングコンテスト2024 のユーザ解説でその存在に言及されています。
floor sum のモノイド積版とは言ったものの、定式化や再帰の細かい部分がちょっとだけオリジナルの floor sum とは異なります（再帰部分についてはだいたい同じにしている人もいるかも？）。 一方、再帰していく流れ自体は floor sum と似ていますし、floor sum（や一般化 floor sum と呼ばれているもの）で計算できる値はこのアルゴリズムで求めることができます。
呼び方を知らなくて、いい呼び方も思い浮かばないので、この記事では floor sum のモノイド積版と呼ぶことにしておきます。
あんまりすごいことは書いてありませんが、ゴールデンウィークに生じたすきまの時間の読み物にどうぞ。
定式化
次の形式の問題が解きたいものです。
問題
$f:\mathbb N\to\mathbb N$ を、非負整数 $a,b$ と正整数 $m\vphantom{A}$ によって $f(x)=\left\lfloor\dfrac{ax+b}m\right\rfloor$ と定める。 モノイド $M\vphantom{A}$ の元 $x,y$ と正整数 $n$ が与えられたとき、次の値を求めよ。
$y ^ {f(0)}xy ^ {f(1)-f(0)}x\cdots xy ^ {f(n)-f(n-1)}$
これのどこが floor sum なんですかというのは元のユーザ解説にある図でうっすら納得しておいてください。
アルゴリズム
再帰的に考えます。 $\operatorname{solve}(n,m,a,b,x,y)$ で上の問題の答えを得ることを考えます。
ベースケースは $f(n)=0$ なるときで、このとき答えは $x ^ {n-1}$ です。
まず、$b\geq m\vphantom a$ のとき、$\operatorname{solve}(n,m,a,b,x,y)=y ^ {\lfloor b/m\rfloor}\operatorname{solve}(n,m,a,b\%m,x,y)$ です。
次に、$a\geq m\vphantom a$ のとき、どのような $t$ についても $f(t)-f(t-1)\gt\left\lfloor\dfrac am\right\rfloor$ が成り立つので、$\operatorname{solve}(n,m,a,b,x,y)=\operatorname{solve}(n,m,a\%m,b,xy ^ {\lfloor a/m\rfloor},y)$ です。
上によって $a\lt m,b\lt m\vphantom a$ とできます。 このとき、求める値は $x ^ {e _ 0}yx ^ {e _ 1}y\cdots yx ^ {e _ k}$ の形になっていることがわかります。 うまく $e _ 0=f ^ \prime(0),e _ 1=f ^ \prime(1)-f ^ \prime(0),\ldots$ となるような $f ^ \prime(x)=\left\lfloor\dfrac{a ^ \prime x+b ^ \prime}{m ^ \prime}\right\rfloor$ が見つかるとうれしいわけですね。
ここで、モノイド $x$ を座標平面上で $x$ 軸正方向に $1$ だけ動く、モノイド $y$ を座標平面状で $y$ 軸正方向に $1$ だけ動く、としたときの動きを図に書いたものを見てみましょう（もとのユーザ解説の図もそのような図になっています）。
この図を、座標平面上の直線 $y=x$ で反転したものを考えます。
[ここには図が入る予定でしたが、保存する前にパソコンが再起動してしまったので、図はありません。各自で図を書いてみてください。]
すると、$f(x)=\left\lceil\dfrac{mx-b}a\right\rceil$ として、ある $k$ が存在して、図が表している値は $y ^ {f(1)}xy ^ {f(2)-f(1)}x\cdots xy ^ {n-f(k)}$ となります。 このことから、$f ^ \prime(x)=\left\lfloor\dfrac{m(x+1)-b+a-1}a\right\rfloor$ とすることで $i=1,2,\ldots,k-1$ については $e _ i=f ^ \prime(i)-f ^ \prime(i-1)$ が成り立つことがわかります（競技プログラマのみなさまなら ceil を floor に書き換えたことはいっぱいあるかと思うので、式変形は細かく追いません）。
ちょっとだけ合っていない部分があり、これは $y ^ {e _ k}$ の部分です。これは仕方がないので、再帰する前に端数を計算しておきます。 以上をまとめて、$a\lt m,b\lt m\vphantom a$ のとき次のように再帰することができます。
\[\operatorname{solve}(n,m,a,b,x,y)=\operatorname{solve}\Biggl(\left\lfloor\dfrac{an+b}m\right\rfloor-1,a,m,m-b+a-1,y,x\Biggr)y x^{\lceil((an+b)\%m)/a\rceil}\]
floor sum の再帰と同じ感じで引数の値が減っていくので、停止性や再帰の回数も同様に議論することができます。
C++ での実装は以下のようになります。
template <std::integral I, class Monoid, std::integral SafeInteger = __uint128_t> Monoid floor_prod(I n, I m, I a, I b, Monoid x, Monoid y) { x *= monoid_pow(y, a / m); a %= m; const auto prefix_prod{monoid_pow(y, b / m)}; b %= m; const auto y_max{static_cast<SafeInteger>(a) * n + b}; if (y_max < m) return prefix_prod * monoid_pow(x, n); return prefix_prod * floor_prod(static_cast<I>(y_max / m) - 1, a, m, m + a - b - 1, y, x) * y * monoid_pow(x, static_cast<I>(y_max % m) / a); }
$a,b$ を $m\vphantom a$ 未満にする操作も再帰の $1$ ステップにまとめてしまっています。 答えの prefix と suffix から決まっていくので、これらを引数に与えると末尾再帰にしたりループに書き換えたりもできます。 お好みで書き換えてしまってもいいでしょう。
モノイドの渡し方を、演算子をオーバーロードしたクラスのインスタンスを渡す形にしています。
他の渡し方にしたいときもお好みで書き換えるとよいでしょう。
monoid_pow 関数も適当に実装しちゃいましょう。
この実装では、$an+b$ を計算する際に $a,n,b$ が入っている型のままではオーバーフローが発生する可能性があるので、より大きい型にキャストしています。
与えられるパラメータが小さいことがわかっている場合や、オーバーフローが生じないよう丁寧な式変形と計算を行う場合は、より大きい型は不要でしょう。
時間計算量について考えます。
モノイド積の計算にかかる時間計算量を $C$ とします。 また、モノイド $a$ と非負整数 $n$ に対して $a ^ n$ を計算するのに $O(\log n)$ 回のモノイド積を行うとします。
monoid_pow 以外は、再帰ごとに定数回の演算です。
monoid_pow(y, a / m) は、$1$ 回の再帰で $O\left(\log\dfrac am\right)$ 回のモノイド積を行います。
ここで、再帰が進むごとに $\dfrac am\vphantom a$ は $\dfrac am,\dfrac m{a\% m},\dfrac{a\% m}{m\%(a\% m)},\cdots$ となるので、全体では $O(\log a)$ 回のモノイド積となります。
同様に、monoid_pow(y, b / m) は全体で $O(\log b)$ 回、monoid_pow(x, static_cast<I>(y_max % m) / a) は全体で $O(\log m)$ 回と評価できます。
よって、全体の時間計算量は $O(C(\log n+\log m+\log a+\log b))$ 時間と評価できます。めでたしめでたし。
verify
この時点でかんたんなモノイドを与え、計算してもらった結果を人間の目で確かめておくことで、安心して次に進めそうですね。
かんたんなモノイドの例としては、文字列の結合（$x=$R$,y=$U など）や、二次元ベクトルの和（$x=(1,0),y=(0,1)$）などを与えるとよいでしょう。
前者は小さめのケース向け、後者は大きめのケース向けです。
前者は自由モノイドなので、これが正しく計算できているならだいたい大丈夫そうですね。
後者はこれに可換性を課したもので、オーバーフロー対策がうまくいっているかなどを確かめるのにちょうどいいかもしれません。
結果が合っていたら安心して次に進みましょう。
floor sum を解く
あとは、個別の問題に対して、floor sum のモノイド積版が求めたい値を返すような適切なモノイドを設計するだけです。
まずは、次の問題が解けるようなモノイドを考えてみましょう。
問題：floor sum
非負整数 $a,b$ と正整数 $n,m\vphantom{A}$ が与えられる。 次の値を求めよ。
$\displaystyle\sum _ {i=0} ^ {n-1}\left\lfloor\dfrac{ai+b}m\right\rfloor$
思いつきましたか？ 例えば、次のようにするとよいです。
- 台集合：非負整数の $3$ つ組 ${\mathbb Z _ {\geq0}} ^ 3$
- 演算：$(a,b,c)\cdot(d,e,f)=(a+d,b+e,c+bd+f)$
- 単位元：$(0,0,0)$
- モノイドの元 $x$：$(1,0,0)$
- モノイドの元 $y$：$(0,1,0)$
気持ちとしては、$($区間内の $x$ の合計$,$ 区間内の $y$ の合計$,$ 区間内で足される値の合計$)$ のようになりますね。
求める値は $3$ つ組の $3$ つめの値として得られます。
実装できたら ACL Practice Contest などに投げ込んでみましょう。 たぶん正解になると思います。
環の値の積の和
次の問題が解けるようなモノイドを考えてみましょう。
問題
非負整数 $a,b$ と正整数 $n,m\vphantom{A}$ 、および環 $R$ の元 $A,B$ が与えられる。 次の値を求めよ。
$\displaystyle\sum _ {i=0} ^ {n-1}A ^ iB ^ {\lfloor(ai+b)/m\rfloor}$
$A,B$ に逆元が存在するとは限らないことや、$A,B$ の積が可換とは限らないことに注意してください。
思いつきましたか？ 例えば、次のようにするとよいです。
- 台集合：$R$ の元の $3$ つ組 $R ^ 3$
- 演算：$(a,b,c)\cdot(d,e,f)=(ad,be,c+afb)$
- 単位元：$(1,1,0)$
- モノイドの元 $x$：$(A,1,1)$
- モノイドの元 $y$：$(1,B,0)$
ただし、ここで $0,1$ はそれぞれ環 $R$ における和・積の単位元とします。 気持ちは floor sum のときとあまり大きく変わっていない気がしますね。
求める値は $3$ つ組の $3$ つめの値として得られます。
実装できたら
などに投げ込んでみましょう。 上の問題では $R$ は $\mathbb F _ {998244353}$ 係数の $N$ 次正方行列に標準的な和と積を入れたものですね。
一般化 floor sum
次の問題を考えましょう。
問題：floor sum
非負整数 $a,b$ と正整数 $n,m\vphantom{A}$ 、$x$ について $p$ 次、$y$ について $q$ 次の整数係数多項式 $f(x,y)$ が与えられる。 次の値を求めよ。
$\displaystyle\sum _ {i=0} ^ {n-1}f\Biggl(i,\left\lfloor\dfrac{ai+b}m\right\rfloor\Biggr)$
天下り的ですが、$\mathbb Q\lbrack\!\lbrack s,t\rbrack\!\rbrack/(s ^ {p+1},t ^ {q+1})$ の元 $A=\mathrm{e} ^ s,B=\mathrm{e} ^ t$ を考えます。
すると、\[ \begin{aligned}&\lbrack s ^ it ^ j\rbrack\sum _ {x=0} ^ {N-1}A ^ xB ^ {\lfloor(ax+b)/m\rfloor}\\ {}={}&\lbrack s ^ it ^ j\rbrack\sum _ {x=0} ^ {N-1}\mathrm{e} ^ {xs}\mathrm{e} ^ {\lfloor(ax+b)/m\rfloor t}\\ {}={}&\sum _ {x=0} ^ {N-1}\lbrack s ^ i\rbrack\mathrm{e} ^ {xs}\lbrack t ^ j\rbrack\mathrm{e} ^ {\lfloor(ax+b)/m\rfloor t}\\ {}={}&\dfrac1{i!j!}\sum _ {x=0} ^ {N-1}x ^ i\left\lfloor\dfrac{ax+b}m\right\rfloor ^ j \end{aligned} \] となるため、$\displaystyle\sum _ {x=0} ^ {N-1}A ^ xB ^ {\lfloor(ax+b)/m\rfloor}$ を求められれば $\Theta(pq)$ 時間で答えを求めることができます。
単に $\mathbb Q\lbrack\!\lbrack s,t\rbrack\!\rbrack/(s ^ {p+1},t ^ {q+1})$ の元どうしの積を求めるのは（$\mathbb Q$ どうしの四則演算を定数時間として） $O\bigl(p ^ 2q ^ 2\bigr)$ 時間などになりますが、管理する $3$ つ組の $1$ つめの値は常に $A$ のべき乗であること（特に、$t$ を含む項が登場しないこと）を利用すると、$1$ つめの値を左からかけることは $O\bigl(p ^ 2q\bigr)$ 時間などでできます。 $2$ つめの値についても同様にして $O\bigl(pq ^ 2\bigr)$ 時間とできます。
ここから、モノイドの元どうしの積が $O\bigl(pq(p+q)\bigr)$ 時間でできます。 多項式どうしの積をフーリエ変換で高速化できるような係数で管理している場合は $O(pq\log pq)$ 時間などにもできそうですね。
これに提出したら TLE しました。いかがでしたか？
非再帰にして $p=1,q=2$ について特殊化すると通りました（$1914\operatorname{ms}$） うーん……
私の実装が悪い気もするので、通った実装を置いておきます。 悪いところを見つけたらこちらまで（画面下部を指で示すぽかいこ）
#include <type_traits> template <typename U, typename... Args> struct has_power_method { private: template <typename V, typename... PowArgs> static auto check(int) -> decltype(std::declval<V>().power(std::declval<PowArgs>()...), std::true_type{}); template <typename, typename...> static auto check(...) -> std::false_type; public: static constexpr bool value = decltype(check<U, Args...>(0))::value; }; #include <concepts> #include <cassert> template <typename Monoid, std::integral I> auto monoid_pow(Monoid x, I exp, Monoid base = Monoid::unit()) { if constexpr (has_power_method<Monoid, I, Monoid>::value) { return x.power(exp, base); } else if constexpr (has_power_method<Monoid, I>::value) { return base * x.power(exp); } else { assert(exp >= 0); Monoid res{base}; while (exp) { if (exp & 1) res *= x; x *= x; exp >>= 1; } return res; } } template <std::integral I, class Monoid, std::integral SafeInteger = __uint128_t> Monoid floor_prod(I n, I m, I a, I b, Monoid x, Monoid y) { Monoid prefix_prod{Monoid::unit()}, suffix_prod{Monoid::unit()}; while (true) { x = monoid_pow(y, a / m, x); a %= m; prefix_prod = monoid_pow(y, b / m, prefix_prod); b %= m; const auto y_max{static_cast<SafeInteger>(a) * n + b}; if (y_max < m) return monoid_pow(x, n, prefix_prod) * suffix_prod; suffix_prod = monoid_pow(x, static_cast<I>(y_max % m) / a, y) * suffix_prod; n = static_cast<I>(y_max / m) - 1; std::swap(a, m); std::swap(x, y); b = m + a - b - 1; } } #include <array> #include <numeric> #include <algorithm> #include <ranges> #include <atcoder/convolution> #include <iostream> namespace pokaiko { template <typename T, std::size_t p, std::size_t q> class floor_sum_polynomial { inline static auto factorial{[] { std::array<T, std::max(p, q) + 1> factorial{}; std::iota(begin(factorial), end(factorial), T{}); factorial.front() = T{1}; std::inclusive_scan(begin(factorial), end(factorial), begin(factorial), std::multiplies{}); return factorial; }()}, inv_factorial{[] { std::array<T, std::max(p, q) + 1> inv_factorial{}; std::iota(begin(inv_factorial), end(inv_factorial), T{1}); inv_factorial.back() = T{1} / factorial.back(); std::inclusive_scan(rbegin(inv_factorial), rend(inv_factorial), rbegin(inv_factorial), std::multiplies{}); return inv_factorial; }()}; public: T A, B; std::array<std::array<T, p + 1>, q + 1> sum; floor_sum_polynomial(T a, T b, std::array<std::array<T, p + 1>, q + 1>&& s) : A(a), B(b), sum(std::move(s)) {} static floor_sum_polynomial x() { std::array<std::array<T, p + 1>, q + 1> sum; for (auto&& row : sum) std::ranges::fill(row, T{}); sum[0][0] = T{1}; return floor_sum_polynomial{T{1}, T{}, std::move(sum)}; } static floor_sum_polynomial y() { std::array<std::array<T, p + 1>, q + 1> sum; for (auto&& row : sum) std::ranges::fill(row, T{}); return floor_sum_polynomial{T{}, T{1}, std::move(sum)}; } static floor_sum_polynomial unit() { std::array<std::array<T, p + 1>, q + 1> sum; for (auto&& row : sum) std::ranges::fill(row, T{}); return floor_sum_polynomial{T{}, T{}, std::move(sum)}; } floor_sum_polynomial& operator*=(floor_sum_polynomial other) { if constexpr (p == 1) { for (auto&& row : other.sum) row[1] += row[0] * A; } else { static std::vector<T> a_coef(p + 1); { T a{1}; for (std::size_t i{}; i <= p; ++i) { a_coef[i] = a * inv_factorial[i]; a *= A; } } for (auto&& row : other.sum) std::ranges::copy(atcoder::convolution(std::vector<T>{begin(row), end(row)}, a_coef) | std::views::take(p + 1), begin(row)); } if constexpr (q == 2) { for (std::size_t i{}; i <= p; ++i) { T s0i = other.sum[0][i], s1i = other.sum[1][i], s2i = other.sum[2][i]; other.sum[2][i] += (other.sum[0][i] * B * ((T::mod() + 1) / 2) + other.sum[1][i]) * B; other.sum[1][i] += other.sum[0][i] * B; } } else { static std::vector<T> b_coef(q + 1), column(q + 1); { T b{1}; for (std::size_t i{}; i <= q; ++i) { b_coef[i] = b * inv_factorial[i]; b *= B; } } for (std::size_t i{}; i <= p; ++i) { std::ranges::copy(other.sum | std::views::transform([i](const auto& row) { return row[i]; }), begin(column)); for (std::size_t j{}; const auto& v : atcoder::convolution(column, b_coef) | std::views::take(q + 1)) other.sum[j++][i] = v; } } for (std::size_t i{}; i <= q; ++i) for (std::size_t j{}; j <= p; ++j) sum[i][j] += other.sum[i][j]; A += other.A; B += other.B; return *this; } floor_sum_polynomial operator*(floor_sum_polynomial const& other) const { return floor_sum_polynomial{*this} *= other; } T ans(const std::array<std::array<T, p + 1>, q + 1>& coef) const { T result{}; for (std::size_t i{}; i <= q; ++i) for (std::size_t j{}; j <= p; ++j) result += sum[i][j] * coef[i][j] * factorial[i] * factorial[j]; return result; } }; } int main() { using namespace std; using modint0 = atcoder::static_modint<998244353>; using modint1 = atcoder::static_modint<1107296257>; unsigned T; cin >> T; const auto solve{ []<class T>(const T&, unsigned n, unsigned m, unsigned a, unsigned b1, unsigned b2) { using monoid = pokaiko::floor_sum_polynomial<T, 1, 2>; T ans{}; array<array<T, 2>, 3> coef1{{{T{b1} * b2, T{a} * (b1 + b2)}, {-(T{b2} * 2 + m) * m, -T{a} * m * 2}, {T{m} * m, 0}}}; array<array<T, 2>, 3> coef2{{{T{b1} * b2, T{a} * (b1 + b2)}, {-(T{b1} * 2 - m) * m, -T{a} * m * 2}, {T{m} * m, 0}}}; ans += floor_prod<unsigned, monoid, unsigned long>(n, m, a, b1, monoid::x(), monoid::y()).ans(coef1); ans += floor_prod<unsigned, monoid, unsigned long>(n, m, a, b2, monoid::x(), monoid::y()).ans(coef2); return ans / 2 + T{a} * a * n * (n - 1) * (2 * n - 1) / 6; } }; const auto coef_1_to_0{modint1{modint0::mod()}.inv()}; const auto coef_0_to_1{modint0{modint1::mod()}.inv()}; for (unsigned i{}; i < T; ++i) [&solve, coef_1_to_0, coef_0_to_1] { unsigned N, M, A, B1, B2; cin >> N >> M >> A >> B1 >> B2; if (B1 < B2) swap(B1, B2); cout << ((solve(modint0{}, N, M, A, B1, B2) * coef_0_to_1).val() * static_cast<unsigned long>(modint1::mod()) + (solve(modint1{}, N, M, A, B1, B2) * coef_1_to_0).val() * static_cast<unsigned long>(modint0::mod())) % (static_cast<unsigned long>(modint0::mod()) * modint1::mod()) << '\n'; }(); return 0; }
$1$ 回計算するだけなら $p=q=256,1\leq n,m,a,b\leq10 ^ 9$ とかでも数秒で計算してくれたので、そんなにめちゃくちゃじゃないと思うんですけど…… どうなんでしょうか
おわり
書き上げてから思ったんですけど、この記事に書いてあることは本家のユーザ解説にもだいたい書いてありませんか？ 行間を埋めているような気もしつつ、この記事にある行間を埋められるならもとのユーザ解説の行間も埋められるような気がします。
がんばれぽかいこ！いつの日か本質的な記事を書き上げる日を夢見て───★
