「二項係数の和（オンライン）」の実装を修正

STYLE_GUIDE.md

@@ -8,6 +8,7 @@
 - 句読点は基本的に `、` と `。` とする（`,` と `.` などでない）。
   - ただし、`Library Checker (Yosupo Judge), yukicoder, AOJ` のように英数字どうしを並べる際のカンマまで避ける必要はない（どちらでもよい）。
 - 複数の語句を並べる際の区切り記号は `、` または `, ` とし、上の規約にも従う（`・` や ` / `、`/` などでない）。
+  - ただし、TeX 数式中の `,` の後ろには空白はあってもなくてもよい。
 - マークダウンの箇条書きのインデントはスペース 2 文字とする。
 - マークダウンの番号なし箇条書きは `- ` とする（`* ` や `  - ` などでない）。
 - マークダウンの番号付き箇条書きは常に `1. ` とする（`2. ` 以降は使用しない）。

docs/math/combinatorics/online-binomial-sum.md

@@ -8,15 +8,18 @@ documentation_of: math/combinatorics/online-binomial-sum.hpp
 - 整数 $n,m$ と重み $r$ に対する二項係数の prefix sum を $\displaystyle F(n,m)=\sum_{i=0}^{n-1}r^i\binom{m}{i}$ とおく。
 - 半開区間の左端、右端をそれぞれ $l,u$ とする。 $\displaystyle \sum_{i=l}^{u-1}r^i\binom{m}{i}$ をオンラインで求める。
 - $\binom{m}{i}=0\;(i>m)$ として扱う。
-- $r=0$ や $r=-1$ でも $r+1$ による除算は行わない。
+- クエリでは前計算を行った最も近い点から復元する。
+- $r=0$ のときは閉形式を用い、 $r=-1$ でも $r+1$ による除算は行わない。
 
 ## 使い方
 
 - `max_m` を $M$ とおく。
 - `OnlineBinomialSum<T>(int max_m, T r = T(1))`
   - $0\le m\le M$ のクエリに対する前計算を行う。
   - `r` は重みで、省略時は $1$ である。
-  - 前提: `T` は四則演算を持つ。`std::numeric_limits<T>::is_integer` が `false` の場合、 $T(1),T(2),\ldots,T(M)$ で除算できる。
+  - 前提: `T` は四則演算と `T()` との等値比較を持つ。
+  - 前提: $r \ne 0$ の場合は `r` で除算できる。
+  - 前提: `std::numeric_limits<T>::is_integer` が `false` の場合、 $T(1),T(2),\ldots,T(M)$ で除算できる。
   - 備考: 整数型では中間値が `T` の範囲を超えない必要がある。
 - `T binom_prefix_sum(int n, int m) const`
   - $\displaystyle \sum_{i=0}^{n-1}r^i\binom{m}{i}$ を返す。

math/combinatorics/online-binomial-sum.hpp

@@ -5,8 +5,10 @@
 // 0 <= l <= u と 0 <= m <= max_m を仮定する。
 // m のバケット境界の累積和を n のバケット境界でサンプルし、
 // バケット内の二項係数を前計算する。
-// T は四則演算を持つ型で、std::numeric_limits<T>::is_integer が
-// false の場合は T(1) から T(max_m) で除算できる。
+// T は四則演算と T() との等値比較を持つ型で、
+// r != T() の場合は r で除算できることを前提とする。
+// std::numeric_limits<T>::is_integer が false の場合は、
+// さらに T(1) から T(max_m) で除算できることを前提とする。
 // 時間計算量は前計算 O(max_m √max_m)、クエリ O(√max_m)。
 // 空間計算量は O(max_m)。
 
@@ -18,16 +20,21 @@ template <class T> struct OnlineBinomialSum {
     int max_m;
     int bucket_size;
     T r;
+    bool r_is_zero;
     std::vector<int> prefix_sum_offset;
     std::vector<T> prefix_sum_table;
     std::vector<T> prefix_term_table;
     std::vector<int> weighted_binomial_offset;
     std::vector<T> weighted_binomial_table;
     std::vector<T> integer_inverse;
+    T r_inverse;
 
-    explicit OnlineBinomialSum(int m, T r = T(1)) : max_m(m), r(r) {
+    explicit OnlineBinomialSum(int m, T r = T(1))
+        : max_m(m), bucket_size(1), r(r), r_is_zero(r == T()) {
         assert(max_m >= 0);
-        bucket_size = 1;
+        if (r_is_zero) {
+            return;
+        }
         while (1LL * bucket_size * bucket_size < max_m + 1) {
             bucket_size <<= 1;
         }
@@ -55,6 +62,7 @@ template <class T> struct OnlineBinomialSum {
         }
 
         if constexpr (!std::numeric_limits<T>::is_integer) {
+            r_inverse = T(1) / r;
             integer_inverse.assign(max_m + 1, T());
             for (int i = 1; i <= max_m; ++i) {
                 integer_inverse[i] = T(1) / T(i);
@@ -64,7 +72,7 @@ template <class T> struct OnlineBinomialSum {
         const int bucket_count = max_m / bucket_size + 1;
         prefix_sum_offset.assign(bucket_count + 1, 0);
         for (int b = 0; b < bucket_count; ++b) {
-            prefix_sum_offset[b + 1] = prefix_sum_offset[b] + b + 1;
+            prefix_sum_offset[b + 1] = prefix_sum_offset[b] + b + 2;
         }
 
         prefix_sum_table.assign(prefix_sum_offset[bucket_count], T());
@@ -92,6 +100,9 @@ template <class T> struct OnlineBinomialSum {
                     }
                 }
             }
+            sum += term;
+            prefix_sum_table[offset + b + 1] = sum;
+            prefix_term_table[offset + b + 1] = T();
         }
     }
 
@@ -105,10 +116,14 @@ template <class T> struct OnlineBinomialSum {
         if (n > m) {
             n = m + 1;
         }
+        if (r_is_zero) {
+            return T(1);
+        }
 
         const int bucket = m / bucket_size;
         const int base = bucket * bucket_size;
         const int d = m - base;
+        const int prefix_offset = prefix_sum_offset[bucket];
         const int weight_offset = weighted_binomial_offset[d];
 
         int last_j = d;
@@ -118,34 +133,109 @@ template <class T> struct OnlineBinomialSum {
         const int first_n = n - last_j;
 
         int sample_index = first_n / bucket_size;
-        if (sample_index > bucket) {
-            sample_index = bucket;
+        if (sample_index > bucket + 1) {
+            sample_index = bucket + 1;
         }
+        int next_sample_index = sample_index + 1;
+        if (next_sample_index > bucket + 1) {
+            next_sample_index = bucket + 1;
+        }
+        const auto restore_cost = [&](int index) -> int {
+            const int sample_n = index * bucket_size;
+            if (sample_n < first_n) {
+                return n - sample_n;
+            }
+            if (sample_n > n) {
+                return sample_n - first_n;
+            }
+            return n - first_n;
+        };
+        if (restore_cost(next_sample_index) < restore_cost(sample_index)) {
+            sample_index = next_sample_index;
+        }
+
         const int sample_n = sample_index * bucket_size;
-        const int sample_offset = prefix_sum_offset[bucket] + sample_index;
+        const int sample_offset = prefix_offset + sample_index;
+        const T base_term = prefix_term_table[prefix_offset + bucket];
+
+        const auto move_right = [&](int current_n, T &sum, T &term) -> void {
+            sum += term;
+            if (current_n < base) {
+                term *= r;
+                term *= T(base - current_n);
+                if constexpr (std::numeric_limits<T>::is_integer) {
+                    term /= T(current_n + 1);
+                } else {
+                    term *= integer_inverse[current_n + 1];
+                }
+            } else {
+                term = T();
+            }
+        };
+        const auto move_left = [&](int current_n, T &sum, T &term) -> void {
+            if (current_n > base + 1) {
+                term = T();
+                return;
+            }
+            if (current_n == base + 1) {
+                term = base_term;
+                sum -= term;
+                return;
+            }
+            if constexpr (std::numeric_limits<T>::is_integer) {
+                term /= r;
+                term *= T(current_n);
+                term /= T(base - current_n + 1);
+            } else {
+                term *= r_inverse;
+                term *= T(current_n);
+                term *= integer_inverse[base - current_n + 1];
+            }
+            sum -= term;
+        };
 
         T sum = prefix_sum_table[sample_offset];
         T term = prefix_term_table[sample_offset];
         T ans = T();
-        for (int current_n = sample_n; current_n <= n; ++current_n) {
-            if (current_n >= first_n) {
+        if (sample_n < first_n) {
+            for (int current_n = sample_n; current_n < first_n; ++current_n) {
+                move_right(current_n, sum, term);
+            }
+            for (int current_n = first_n; current_n <= n; ++current_n) {
                 const int j = n - current_n;
                 ans += weighted_binomial_table[weight_offset + j] * sum;
+                if (current_n < n) {
+                    move_right(current_n, sum, term);
+                }
             }
-            if (current_n < n) {
-                sum += term;
-                if (current_n < base) {
-                    term *= r;
-                    term *= T(base - current_n);
-                    if constexpr (std::numeric_limits<T>::is_integer) {
-                        term /= T(current_n + 1);
-                    } else {
-                        term *= integer_inverse[current_n + 1];
-                    }
-                } else {
-                    term = T();
+        } else if (sample_n > n) {
+            for (int current_n = sample_n; current_n > n; --current_n) {
+                move_left(current_n, sum, term);
+            }
+            for (int current_n = n; current_n >= first_n; --current_n) {
+                const int j = n - current_n;
+                ans += weighted_binomial_table[weight_offset + j] * sum;
+                if (current_n > first_n) {
+                    move_left(current_n, sum, term);
+                }
+            }
+        } else {
+            T right_sum = sum;
+            T right_term = term;
+            for (int current_n = sample_n; current_n <= n; ++current_n) {
+                const int j = n - current_n;
+                ans += weighted_binomial_table[weight_offset + j] * right_sum;
+                if (current_n < n) {
+                    move_right(current_n, right_sum, right_term);
                 }
             }
+            T left_sum = sum;
+            T left_term = term;
+            for (int current_n = sample_n; current_n > first_n; --current_n) {
+                move_left(current_n, left_sum, left_term);
+                const int j = n - current_n + 1;
+                ans += weighted_binomial_table[weight_offset + j] * left_sum;
+            }
         }
         return ans;
     }