mod/floor_sum_of_linear_polynomial.hpp

View this file on GitHub
Last update: 2024-01-23 03:59:43+09:00
Include: #include "mod/floor_sum_of_linear_polynomial.hpp"

Depends on

Verified with

Code

#include "mod/floor_monoid_product.hpp"
#include "alg/monoid/monoid_for_floor_sum.hpp"

// sum(i^k*floor) for k=0,1,...,K
template <typename T, int K, typename U>
array<T, K + 1> floor_sum_of_linear_polynomial(U N, U a, U b, U mod) {
  static_assert(is_same_v<U, u64> || is_same_v<U, u128>);
  assert(a == 0 || N < (U(-1) - b) / a);
  using Mono = Monoid_for_floor_sum<T, K>;
  auto z = floor_monoid_product<Mono>(Mono::to_x(), Mono::to_y(), N, a, b, mod);
  return z.dp[1];
};

#line 1 "mod/floor_sum_of_linear_polynomial.hpp"

#line 2 "alg/monoid_pow.hpp"

// chat gpt
template <typename U, typename Arg1, typename Arg2>
struct has_power_method {
private:
  // ヘルパー関数の実装
  template <typename V, typename A1, typename A2>
  static auto check(int)
      -> decltype(std::declval<V>().power(std::declval<A1>(),
                                          std::declval<A2>()),
                  std::true_type{});
  template <typename, typename, typename>
  static auto check(...) -> std::false_type;

public:
  // メソッドの有無を表す型
  static constexpr bool value = decltype(check<U, Arg1, Arg2>(0))::value;
};

template <typename Monoid>
typename Monoid::X monoid_pow(typename Monoid::X x, ll exp) {
  using X = typename Monoid::X;
  if constexpr (has_power_method<Monoid, X, ll>::value) {
    return Monoid::power(x, exp);
  } else {
    assert(exp >= 0);
    X res = Monoid::unit();
    while (exp) {
      if (exp & 1) res = Monoid::op(res, x);
      x = Monoid::op(x, x);
      exp >>= 1;
    }
    return res;
  }
}
#line 2 "mod/floor_monoid_product.hpp"

// https://yukicoder.me/submissions/883884
// https://qoj.ac/contest/1411/problem/7620
// U は範囲内で ax+b がオーバーフローしない程度
// yyy x yyyy x ... yyy x yyy (x を N 個)
// k 個目の x までに floor(ak+b,m) 個の y がある
// my<=ax+b における lattice path における辺の列と見なせる
template <typename Monoid, typename X, typename U>
X floor_monoid_product(X x, X y, U N, U a, U b, U m) {
  U c = (a * N + b) / m;
  X pre = Monoid::unit(), suf = Monoid::unit();
  while (1) {
    const U p = a / m, q = b / m;
    a %= m, b %= m;
    x = Monoid::op(x, monoid_pow<Monoid>(y, p));
    pre = Monoid::op(pre, monoid_pow<Monoid>(y, q));
    c -= (p * N + q);
    if (c == 0) break;
    const U d = (m * c - b - 1) / a + 1;
    suf = Monoid::op(y, Monoid::op(monoid_pow<Monoid>(x, N - d), suf));
    b = m - b - 1 + a, N = c - 1, c = d;
    swap(m, a), swap(x, y);
  }
  x = monoid_pow<Monoid>(x, N);
  return Monoid::op(Monoid::op(pre, x), suf);
}
#line 1 "alg/monoid/monoid_for_floor_sum.hpp"

// sum i^kfloor: floor path で (x,y) から x 方向に進むときに x^ky を足す
template <typename T, int K>
struct Monoid_for_floor_sum {
  using ARR = array<array<T, K + 1>, 2>;
  struct Data {
    ARR dp;
    T dx, dy;
  };

  using value_type = Data;
  using X = value_type;
  static X op(X a, X b) {
    static T comb[K + 1][K + 1];
    if (comb[0][0] != T(1)) {
      comb[0][0] = T(1);
      FOR(i, K) FOR(j, i + 1) {
        comb[i + 1][j] += comb[i][j], comb[i + 1][j + 1] += comb[i][j];
      }
    }
    T pow = 1;
    FOR(j, K + 1) {
      FOR(i, K - j + 1) {
        T x = comb[i + j][i];
        a.dp[0][i + j] += b.dp[0][i] * pow * x;
        a.dp[1][i + j] += (b.dp[0][i] * a.dy + b.dp[1][i]) * pow * x;
      }
      pow *= a.dx;
    }
    a.dx += b.dx, a.dy += b.dy;
    return a;
  }

  static X to_x() {
    X x = unit();
    x.dp[0][0] = 1, x.dx = 1;
    return x;
  }
  static X to_y() {
    X x = unit();
    x.dy = 1;
    return x;
  }
  static constexpr X unit() { return {ARR{}, T(0), T(0)}; }
  static constexpr bool commute = 0;
};
#line 4 "mod/floor_sum_of_linear_polynomial.hpp"

// sum(i^k*floor) for k=0,1,...,K
template <typename T, int K, typename U>
array<T, K + 1> floor_sum_of_linear_polynomial(U N, U a, U b, U mod) {
  static_assert(is_same_v<U, u64> || is_same_v<U, u128>);
  assert(a == 0 || N < (U(-1) - b) / a);
  using Mono = Monoid_for_floor_sum<T, K>;
  auto z = floor_monoid_product<Mono>(Mono::to_x(), Mono::to_y(), N, a, b, mod);
  return z.dp[1];
};