mod/floor_sum_of_linear_polynomial.hpp

• View this file on GitHub
• Last update: 2024-08-26 01:20:39+09:00
• Include: #include "mod/floor_sum_of_linear_polynomial.hpp"

Depends on

• alg/monoid/monoid_for_floor_sum.hpp
• alg/monoid_pow.hpp
• mod/floor_monoid_product.hpp

Required by

• convex/count_lattice_point_in_convex_polygon_polynomial.hpp

Verified with

• test/1_mytest/floor_sum_of_polynomial.test.cpp
• test/2_library_checker/number_theory/sum_of_floor_of_linear_2.test.cpp
• test/3_yukicoder/2362.test.cpp

Code

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

// 全部非負, T は答, U は ax+b がオーバーフローしない
template <typename T, int K1, int K2, typename U>
array<array<T, K2 + 1>, K1 + 1> floor_sum_of_linear_polynomial_nonnegative(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, K1, K2>;
  auto x = floor_monoid_product<Mono>(Mono::to_x(), Mono::to_y(), N, a, b, mod);
  return x.dp;
};

// sum_{L<=x<R} x^i floor(ax+b,mod)^j
// a+bx が I, U でオーバーフローしない
template <typename T, int K1, int K2, typename I>
array<array<T, K2 + 1>, K1 + 1> floor_sum_of_linear_polynomial(I L, I R, I a, I b, I mod) {
  static_assert(is_same_v<I, ll> || is_same_v<I, i128>);
  assert(L <= R && mod > 0);
  if (a < 0) {
    auto ANS = floor_sum_of_linear_polynomial<T, K1, K2, I>(-R + 1, -L + 1, -a, b, mod);
    FOR(i, K1 + 1) {
      if (i % 2 == 1) { FOR(j, K2 + 1) ANS[i][j] = -ANS[i][j]; }
    }
    return ANS;
  }
  assert(a >= 0);
  I ADD_X = L;
  I N = R - L;
  b += a * L;
  I ADD_Y = floor<I>(b, mod);
  b -= ADD_Y * mod;
  assert(a >= 0 && b >= 0);

  using Mono = Monoid_for_floor_sum<T, K1, K2>;
  using Data = typename Mono::Data;
  using U = std::conditional_t<is_same_v<I, ll>, i128, u128>;
  Data A = floor_monoid_product<Mono, Data, U>(Mono::to_x(), Mono::to_y(), N, a, b, mod);
  Data offset = Mono::unit();
  offset.dx = T(ADD_X), offset.dy = T(ADD_Y);
  A = Mono::op(offset, A);
  return A.dp;
};

#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^k1floor^k2: floor path で (x,y) から x 方向に進むときに x^k1y^k2 を足す
template <typename T, int K1, int K2>
struct Monoid_for_floor_sum {
  using ARR = array<array<T, K2 + 1>, K1 + 1>;
  struct Data {
    ARR dp;
    T dx, dy;
  };

  using value_type = Data;
  using X = value_type;
  static X op(X a, X b) {
    static constexpr int n = max(K1, K2);
    static T comb[n + 1][n + 1];
    if (comb[0][0] != T(1)) {
      comb[0][0] = T(1);
      FOR(i, n) FOR(j, i + 1) { comb[i + 1][j] += comb[i][j], comb[i + 1][j + 1] += comb[i][j]; }
    }

    array<T, K1 + 1> pow_x;
    array<T, K2 + 1> pow_y;
    pow_x[0] = 1, pow_y[0] = 1;
    FOR(i, K1) pow_x[i + 1] = pow_x[i] * a.dx;
    FOR(i, K2) pow_y[i + 1] = pow_y[i] * a.dy;

    // +dy
    FOR(i, K1 + 1) {
      FOR_R(j, K2 + 1) {
        T x = b.dp[i][j];
        FOR(k, j + 1, K2 + 1) b.dp[i][k] += comb[k][j] * pow_y[k - j] * x;
      }
    }
    // +dx
    FOR(j, K2 + 1) {
      FOR_R(i, K1 + 1) { FOR(k, i, K1 + 1) a.dp[k][j] += comb[k][i] * pow_x[k - i] * b.dp[i][j]; }
    }

    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"

// 全部非負, T は答, U は ax+b がオーバーフローしない
template <typename T, int K1, int K2, typename U>
array<array<T, K2 + 1>, K1 + 1> floor_sum_of_linear_polynomial_nonnegative(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, K1, K2>;
  auto x = floor_monoid_product<Mono>(Mono::to_x(), Mono::to_y(), N, a, b, mod);
  return x.dp;
};

// sum_{L<=x<R} x^i floor(ax+b,mod)^j
// a+bx が I, U でオーバーフローしない
template <typename T, int K1, int K2, typename I>
array<array<T, K2 + 1>, K1 + 1> floor_sum_of_linear_polynomial(I L, I R, I a, I b, I mod) {
  static_assert(is_same_v<I, ll> || is_same_v<I, i128>);
  assert(L <= R && mod > 0);
  if (a < 0) {
    auto ANS = floor_sum_of_linear_polynomial<T, K1, K2, I>(-R + 1, -L + 1, -a, b, mod);
    FOR(i, K1 + 1) {
      if (i % 2 == 1) { FOR(j, K2 + 1) ANS[i][j] = -ANS[i][j]; }
    }
    return ANS;
  }
  assert(a >= 0);
  I ADD_X = L;
  I N = R - L;
  b += a * L;
  I ADD_Y = floor<I>(b, mod);
  b -= ADD_Y * mod;
  assert(a >= 0 && b >= 0);

  using Mono = Monoid_for_floor_sum<T, K1, K2>;
  using Data = typename Mono::Data;
  using U = std::conditional_t<is_same_v<I, ll>, i128, u128>;
  Data A = floor_monoid_product<Mono, Data, U>(Mono::to_x(), Mono::to_y(), N, a, b, mod);
  Data offset = Mono::unit();
  offset.dx = T(ADD_X), offset.dy = T(ADD_Y);
  A = Mono::op(offset, A);
  return A.dp;
};

Back to top page