nt/sigma_0_sum.hpp

• View this file on GitHub
• Last update: 2024-08-29 05:12:14+09:00
• Include: #include "nt/sigma_0_sum.hpp"
• Link: https://oeis.org/A006218

Depends on

• nt/convex_floor_sum.hpp

Verified with

• test/1_mytest/sigma_0_sum.test.cpp
• test/5_atcoder/abc230e.test.cpp

Code

#include "nt/convex_floor_sum.hpp"

// sum_[1,N] sigma_0(n)
template <typename T = u64>
T sigma_0_sum_small(u64 N) {
  u32 sq = sqrtl(N);
  T ANS = 0;
  for (u32 d = 1; d <= sq; ++d) { ANS += N / d; }
  return 2 * ANS - u64(sq) * sq;
}

// https://oeis.org/A006218
// sigma0(1)+...+sigma0(N) = sum floor(N/i)
template <typename T = u64>
T sigma_0_sum_large(u64 N) {
  u32 sq = sqrtl(N);
  auto above = [&](u128 x, u128 y) -> bool { return y * (sq - x) > N; };
  auto slope = [&](u128 x, u128 a, u128 b) -> bool {
    x = sq - x;
    return a * x * x <= N * b;
  };
  T ANS = convex_floor_sum<u64, T>(sq, above, slope);
  return 2 * ANS - u64(sq) * sq;
}

template <typename T = u64>
T sigma_0_sum(u64 N) {
  return (N < (1e14) ? sigma_0_sum_small<T>(N) : sigma_0_sum_large<T>(N));
}

#line 1 "nt/convex_floor_sum.hpp"

// f: 凸, 非負, 単調増加を仮定
// above(x,y) : y > f(x)
// slope(x,a,b) : f'(x) >= a/b
// return : sum_[0,N) floor(f(x))
// https://qoj.ac/contest/1195/problem/6188
template <typename U, typename ANS_TYPE, typename F1, typename F2>
ANS_TYPE convex_floor_sum(U N, F1 above, F2 slope) {
  if (N == 0) return 0;
  auto check = [&](U x, U y) -> bool { return x < N && above(x, y); };

  using T = ANS_TYPE;
  auto max_add = [&](auto f, U& a, U& b, U c, U d) -> void {
    auto dfs = [&](auto& dfs, U c, U d) -> void {
      if (!f(a + c, b + d)) return;
      a += c, b += d, dfs(dfs, c + c, d + d);
      if (f(a + c, b + d)) a += c, b += d;
    };
    dfs(dfs, c, d);
  };

  assert(!above(0, 0));
  U x = 0, y = 0;
  max_add([&](U x, U y) -> bool { return !above(x, y); }, x, y, 0, 1);
  ++y;
  T ANS = 2 * (y - 1); //  [0,1) x [1,y)

  auto add_ANS = [&](U& x, U& y, U a, U b) -> void {
    U x0 = x, y0 = y;
    max_add(check, x, y, a, b);
    U n = (x - x0) / a;
    //  (x0,y0) to (x,y)
    ANS += 2 * (y0 - 1) * (x - x0);
    ANS += (x - x0 + 1) * (y - y0 + 1) - (n + 1);
  };

  add_ANS(x, y, 1, 0);
  vc<tuple<U, U, U, U>> SBT;
  SBT.eb(1, 0, 0, 1);
  while (x < N - 1) {
    U a, b, c, d;
    tie(a, b, c, d) = SBT.back();
    if (!check(x + c, y + d)) {
      POP(SBT);
      continue;
    }
    auto f = [&](u64 a, u64 b) -> bool {
      if (x + a >= N) return 0;
      if (above(x + a, y + b)) return 0;
      if (slope(x + a, d, c)) return 0;
      return 1;
    };
    max_add(f, a, b, c, d);
    if (check(x + a + c, y + b + d)) {
      max_add([&](U a, U b) -> bool { return check(x + a, y + b); }, c, d, a, b);
      SBT.eb(a, b, c, d);
      continue;
    }
    add_ANS(x, y, c, d);
  }
  ANS /= T(2);
  return ANS;
}
#line 2 "nt/sigma_0_sum.hpp"

// sum_[1,N] sigma_0(n)
template <typename T = u64>
T sigma_0_sum_small(u64 N) {
  u32 sq = sqrtl(N);
  T ANS = 0;
  for (u32 d = 1; d <= sq; ++d) { ANS += N / d; }
  return 2 * ANS - u64(sq) * sq;
}

// https://oeis.org/A006218
// sigma0(1)+...+sigma0(N) = sum floor(N/i)
template <typename T = u64>
T sigma_0_sum_large(u64 N) {
  u32 sq = sqrtl(N);
  auto above = [&](u128 x, u128 y) -> bool { return y * (sq - x) > N; };
  auto slope = [&](u128 x, u128 a, u128 b) -> bool {
    x = sq - x;
    return a * x * x <= N * b;
  };
  T ANS = convex_floor_sum<u64, T>(sq, above, slope);
  return 2 * ANS - u64(sq) * sq;
}

template <typename T = u64>
T sigma_0_sum(u64 N) {
  return (N < (1e14) ? sigma_0_sum_small<T>(N) : sigma_0_sum_large<T>(N));
}

Back to top page