Skip to content

dirichlet_prefix.hpp

SECTIONMath INCLUDEnoya/dirichlet_prefix.hpp

Return Q_n={floor(n/i)} in increasing order. Consecutive equal quotients are represented once.

Verified by dirichlet_convolution_and_prefix_sums, dirichlet_inverse_and_prefix_sums.

\[ \displaystyle q_i=\lfloor\frac{n}{i}\rfloor \]

Implementation

View on GitHub

#ifndef NOYA_DIRICHLET_PREFIX_HPP
#define NOYA_DIRICHLET_PREFIX_HPP 1

/// @complexity Time: O(n^(2/3)) arithmetic operations per convolution or
/// division. Space: O(sqrt(n)).

#include "noya/integer_kth_root.hpp"

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <vector>

namespace noya {

namespace dirichlet_prefix_internal {

using u64 = std::uint64_t;
using u128 = unsigned __int128;

inline u64 integer_square_root(u64 value) {
  u64 root = u64(std::sqrt(static_cast<long double>(value)));
  while (u128(root + 1) * (root + 1) <= value) {
    root++;
  }
  while (u128(root) * root > value) {
    root--;
  }
  return root;
}

template <class T> class quotient_table {
public:
  explicit quotient_table(u64 limit) : limit_(limit) {
    assert(limit_ >= 1);
    square_root_ = integer_square_root(limit_);
    small_limit_ =
        u128(square_root_) * square_root_ + square_root_ <= limit_
            ? square_root_
            : square_root_ - 1;
    table_size_ = int(small_limit_ + square_root_ + 1);
    values_.resize(table_size_);
  }

  quotient_table(u64 limit, const std::vector<T> &prefix_values)
      : quotient_table(limit) {
    assert(int(prefix_values.size()) + 1 == table_size_);
    std::copy(prefix_values.begin(), prefix_values.end(), values_.begin() + 1);
  }

  int index_of(u64 value) const {
    assert(1 <= value && value <= limit_);
    int index = value <= small_limit_
                    ? int(value)
                    : table_size_ - int(limit_ / value);
    assert(1 <= index && index < table_size_);
    return index;
  }

  u64 argument(int index) const {
    assert(1 <= index && index < table_size_);
    return u64(index) <= small_limit_
               ? u64(index)
               : limit_ / u64(table_size_ - index);
  }

  std::vector<T> export_values() const {
    return std::vector<T>(values_.begin() + 1, values_.end());
  }

  u64 limit() const { return limit_; }
  u64 square_root() const { return square_root_; }
  u64 small_limit() const { return small_limit_; }
  int table_size() const { return table_size_; }

  T &operator[](int index) { return values_[index]; }
  const T &operator[](int index) const { return values_[index]; }

private:
  u64 limit_ = 0;
  u64 square_root_ = 0;
  u64 small_limit_ = 0;
  int table_size_ = 0;
  std::vector<T> values_;
};

template <class T>
quotient_table<T> convolution(const quotient_table<T> &first,
                              const quotient_table<T> &second) {
  assert(first.limit() == second.limit());
  u64 limit = first.limit();
  quotient_table<T> result(limit);
  int table_size = first.table_size();
  if (limit == 1) {
    result[1] = first[1] * second[1];
    return result;
  }

  std::vector<T> first_value(table_size), second_value(table_size);
  for (int index = 1; index < table_size; index++) {
    first_value[index] = first[index] - first[index - 1];
    second_value[index] = second[index] - second[index - 1];
  }

  u64 cube_root = integer_kth_root(limit, 3);
  u64 split = cube_root * cube_root;
  u64 small_limit = first.small_limit();
  u64 square_root = first.square_root();

  for (u64 left = 1; left <= cube_root; left++) {
    result[first.index_of(left * left)] +=
        first_value[left] * second_value[left];

    if (left * (left + 1) <= small_limit) {
      u64 upper = small_limit / left;
      for (u64 right = left + 1; right <= upper; right++) {
        result[int(left * right)] +=
            first_value[left] * second_value[right] +
            first_value[right] * second_value[left];
      }
    }

    u64 upper = std::min(split / left, small_limit);
    for (u64 right = std::max(left, small_limit / left) + 1;
         right <= upper; right++) {
      result[first.index_of(left * right)] +=
          first_value[left] * second_value[right] +
          first_value[right] * second_value[left];
    }

    if (limit / square_root <= split / left) {
      u64 lower = limit / (split / left + 1) + 1;
      for (u64 block = lower; block <= square_root; block++) {
        int large_index = table_size - int(block);
        result[first.index_of(left * (limit / block))] +=
            first_value[left] * second_value[large_index] +
            first_value[large_index] * second_value[left];
      }
    }
  }

  for (int index = 1; index < table_size; index++) {
    result[index] += result[index - 1];
  }
  for (u64 block = 1; limit / block > split; block++) {
    u64 maximum = limit / block;
    u64 root = integer_square_root(maximum);
    int result_index = table_size - int(block);
    result[result_index] = 0;
    for (u64 left = 1; left <= root; left++) {
      int right_index = first.index_of(maximum / left);
      result[result_index] +=
          first_value[left] * second[right_index] +
          second_value[left] * first[right_index];
    }
    result[result_index] -= first[root] * second[root];
  }
  return result;
}

template <class T>
quotient_table<T> divide(const quotient_table<T> &numerator,
                         quotient_table<T> denominator) {
  assert(numerator.limit() == denominator.limit());
  assert(denominator[1] != T(0));
  u64 limit = numerator.limit();
  quotient_table<T> quotient(limit);
  int table_size = denominator.table_size();
  if (limit == 1) {
    quotient[1] = numerator[1] / denominator[1];
    return quotient;
  }

  T inverse_constant = T(1) / denominator[1];
  for (int index = 0; index < table_size; index++) {
    denominator[index] *= inverse_constant;
  }

  std::vector<T> denominator_value(table_size), quotient_value(table_size),
      remaining_value(table_size);
  for (int index = 1; index < table_size; index++) {
    denominator_value[index] = denominator[index] - denominator[index - 1];
    remaining_value[index] = numerator[index] - numerator[index - 1];
  }

  u64 cube_root = integer_kth_root(limit, 3);
  u64 split = std::max(denominator.square_root(), cube_root * cube_root);
  quotient_value[1] = numerator[1];
  for (int index = 2; index < table_size; index++) {
    u64 argument = denominator.argument(index);
    if (argument > split) {
      break;
    }
    quotient_value[index] =
        remaining_value[index] - quotient_value[1] * denominator_value[index];
    if (argument * argument <= split) {
      remaining_value[denominator.index_of(argument * argument)] -=
          denominator_value[index] * quotient_value[index];
    }
    u64 upper = std::min<u64>(index - 1, split / argument);
    for (u64 other = 2; other <= upper; other++) {
      remaining_value[denominator.index_of(argument * other)] -=
          denominator_value[index] * quotient_value[other] +
          denominator_value[other] * quotient_value[index];
    }
  }

  for (int index = 1; index < table_size; index++) {
    quotient[index] = quotient[index - 1] + quotient_value[index];
  }
  for (u64 block = limit / (split + 1); block > 0; block--) {
    int result_index = table_size - int(block);
    u64 maximum = limit / block;
    u64 root = integer_square_root(maximum);
    quotient[result_index] =
        numerator[result_index] - quotient_value[1] * denominator[result_index] +
        denominator[root] * quotient[root];
    for (u64 left = 2; left <= root; left++) {
      int right_index = denominator.index_of(maximum / left);
      quotient[result_index] -=
          denominator_value[left] * quotient[right_index] +
          quotient_value[left] * denominator[right_index];
    }
  }
  for (int index = 0; index < table_size; index++) {
    quotient[index] *= inverse_constant;
  }
  return quotient;
}

} // namespace dirichlet_prefix_internal

/// @brief Return Q_n={floor(n/i)} in increasing order. Consecutive equal
/// quotients are represented once.
inline std::vector<std::uint64_t> dirichlet_quotients(std::uint64_t limit) {
  dirichlet_prefix_internal::quotient_table<int> layout(limit);
  std::vector<std::uint64_t> result;
  result.reserve(layout.table_size() - 1);
  for (int index = 1; index < layout.table_size(); index++) {
    result.push_back(layout.argument(index));
  }
  return result;
}

/// @brief Given prefix sums of f and g at all increasing values in Q_n,
/// return the corresponding prefix sums of their Dirichlet convolution. Small
/// products are enumerated once around n^(1/3); for a large quotient x, the
/// hyperbola sum groups all equal floor(x/i) values through the Q_n index.
template <class T>
std::vector<T> dirichlet_convolution_prefix_sums(
    std::uint64_t limit, const std::vector<T> &first,
    const std::vector<T> &second) {
  using dirichlet_prefix_internal::convolution;
  using dirichlet_prefix_internal::quotient_table;
  return convolution(quotient_table<T>(limit, first),
                     quotient_table<T>(limit, second))
      .export_values();
}

/// @brief Given prefix sums of f at all increasing values in Q_n, return the
/// prefix sums of its Dirichlet inverse. The recurrence f*g=delta is solved for
/// all small arguments first; each remaining large quotient is then recovered
/// by one grouped divisor-hyperbola equation using those known values.
template <class T>
std::vector<T> dirichlet_inverse_prefix_sums(std::uint64_t limit,
                                             const std::vector<T> &function) {
  using dirichlet_prefix_internal::divide;
  using dirichlet_prefix_internal::quotient_table;
  quotient_table<T> identity(limit);
  for (int index = 1; index < identity.table_size(); index++) {
    identity[index] = T(1);
  }
  return divide(identity, quotient_table<T>(limit, function)).export_values();
}

} // namespace noya

#endif // NOYA_DIRICHLET_PREFIX_HPP
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <vector>

/// @complexity Time: O(n^(2/3)) arithmetic operations per convolution or
/// division. Space: O(sqrt(n)).

/// @complexity Time: O(k ceil(64/k)) = O(64) bounded multiplications for a
/// 64-bit input.  Space: O(1).

namespace noya {

/// @brief Return floor(value^(1/exponent)) for an unsigned 64-bit integer.
/// Binary search uses an exponent-dependent upper bound.  The comparison
/// multiplies only while the next factor is at most value/current, so it is
/// exact and never relies on floating-point rounding or overflowing products.
inline std::uint64_t integer_kth_root(std::uint64_t value, int exponent) {
  assert(1 <= exponent && exponent <= 64);
  if (exponent == 1 || value <= 1) {
    return value;
  }
  auto power_at_most = [&](std::uint64_t base) {
    std::uint64_t product = 1;
    for (int count = 0; count < exponent; count++) {
      if (base != 0 && product > value / base) {
        return false;
      }
      product *= base;
    }
    return true;
  };
  int upper_bit = (64 + exponent - 1) / exponent;
  std::uint64_t low = 0;
  std::uint64_t high = std::uint64_t(1) << upper_bit;
  while (high - low > 1) {
    std::uint64_t middle = low + (high - low) / 2;
    (power_at_most(middle) ? low : high) = middle;
  }
  return low;
}

} // namespace noya

namespace noya {

namespace dirichlet_prefix_internal {

using u64 = std::uint64_t;
using u128 = unsigned __int128;

inline u64 integer_square_root(u64 value) {
  u64 root = u64(std::sqrt(static_cast<long double>(value)));
  while (u128(root + 1) * (root + 1) <= value) {
    root++;
  }
  while (u128(root) * root > value) {
    root--;
  }
  return root;
}

template <class T> class quotient_table {
public:
  explicit quotient_table(u64 limit) : limit_(limit) {
    assert(limit_ >= 1);
    square_root_ = integer_square_root(limit_);
    small_limit_ =
        u128(square_root_) * square_root_ + square_root_ <= limit_
            ? square_root_
            : square_root_ - 1;
    table_size_ = int(small_limit_ + square_root_ + 1);
    values_.resize(table_size_);
  }

  quotient_table(u64 limit, const std::vector<T> &prefix_values)
      : quotient_table(limit) {
    assert(int(prefix_values.size()) + 1 == table_size_);
    std::copy(prefix_values.begin(), prefix_values.end(), values_.begin() + 1);
  }

  int index_of(u64 value) const {
    assert(1 <= value && value <= limit_);
    int index = value <= small_limit_
                    ? int(value)
                    : table_size_ - int(limit_ / value);
    assert(1 <= index && index < table_size_);
    return index;
  }

  u64 argument(int index) const {
    assert(1 <= index && index < table_size_);
    return u64(index) <= small_limit_
               ? u64(index)
               : limit_ / u64(table_size_ - index);
  }

  std::vector<T> export_values() const {
    return std::vector<T>(values_.begin() + 1, values_.end());
  }

  u64 limit() const { return limit_; }
  u64 square_root() const { return square_root_; }
  u64 small_limit() const { return small_limit_; }
  int table_size() const { return table_size_; }

  T &operator[](int index) { return values_[index]; }
  const T &operator[](int index) const { return values_[index]; }

private:
  u64 limit_ = 0;
  u64 square_root_ = 0;
  u64 small_limit_ = 0;
  int table_size_ = 0;
  std::vector<T> values_;
};

template <class T>
quotient_table<T> convolution(const quotient_table<T> &first,
                              const quotient_table<T> &second) {
  assert(first.limit() == second.limit());
  u64 limit = first.limit();
  quotient_table<T> result(limit);
  int table_size = first.table_size();
  if (limit == 1) {
    result[1] = first[1] * second[1];
    return result;
  }

  std::vector<T> first_value(table_size), second_value(table_size);
  for (int index = 1; index < table_size; index++) {
    first_value[index] = first[index] - first[index - 1];
    second_value[index] = second[index] - second[index - 1];
  }

  u64 cube_root = integer_kth_root(limit, 3);
  u64 split = cube_root * cube_root;
  u64 small_limit = first.small_limit();
  u64 square_root = first.square_root();

  for (u64 left = 1; left <= cube_root; left++) {
    result[first.index_of(left * left)] +=
        first_value[left] * second_value[left];

    if (left * (left + 1) <= small_limit) {
      u64 upper = small_limit / left;
      for (u64 right = left + 1; right <= upper; right++) {
        result[int(left * right)] +=
            first_value[left] * second_value[right] +
            first_value[right] * second_value[left];
      }
    }

    u64 upper = std::min(split / left, small_limit);
    for (u64 right = std::max(left, small_limit / left) + 1;
         right <= upper; right++) {
      result[first.index_of(left * right)] +=
          first_value[left] * second_value[right] +
          first_value[right] * second_value[left];
    }

    if (limit / square_root <= split / left) {
      u64 lower = limit / (split / left + 1) + 1;
      for (u64 block = lower; block <= square_root; block++) {
        int large_index = table_size - int(block);
        result[first.index_of(left * (limit / block))] +=
            first_value[left] * second_value[large_index] +
            first_value[large_index] * second_value[left];
      }
    }
  }

  for (int index = 1; index < table_size; index++) {
    result[index] += result[index - 1];
  }
  for (u64 block = 1; limit / block > split; block++) {
    u64 maximum = limit / block;
    u64 root = integer_square_root(maximum);
    int result_index = table_size - int(block);
    result[result_index] = 0;
    for (u64 left = 1; left <= root; left++) {
      int right_index = first.index_of(maximum / left);
      result[result_index] +=
          first_value[left] * second[right_index] +
          second_value[left] * first[right_index];
    }
    result[result_index] -= first[root] * second[root];
  }
  return result;
}

template <class T>
quotient_table<T> divide(const quotient_table<T> &numerator,
                         quotient_table<T> denominator) {
  assert(numerator.limit() == denominator.limit());
  assert(denominator[1] != T(0));
  u64 limit = numerator.limit();
  quotient_table<T> quotient(limit);
  int table_size = denominator.table_size();
  if (limit == 1) {
    quotient[1] = numerator[1] / denominator[1];
    return quotient;
  }

  T inverse_constant = T(1) / denominator[1];
  for (int index = 0; index < table_size; index++) {
    denominator[index] *= inverse_constant;
  }

  std::vector<T> denominator_value(table_size), quotient_value(table_size),
      remaining_value(table_size);
  for (int index = 1; index < table_size; index++) {
    denominator_value[index] = denominator[index] - denominator[index - 1];
    remaining_value[index] = numerator[index] - numerator[index - 1];
  }

  u64 cube_root = integer_kth_root(limit, 3);
  u64 split = std::max(denominator.square_root(), cube_root * cube_root);
  quotient_value[1] = numerator[1];
  for (int index = 2; index < table_size; index++) {
    u64 argument = denominator.argument(index);
    if (argument > split) {
      break;
    }
    quotient_value[index] =
        remaining_value[index] - quotient_value[1] * denominator_value[index];
    if (argument * argument <= split) {
      remaining_value[denominator.index_of(argument * argument)] -=
          denominator_value[index] * quotient_value[index];
    }
    u64 upper = std::min<u64>(index - 1, split / argument);
    for (u64 other = 2; other <= upper; other++) {
      remaining_value[denominator.index_of(argument * other)] -=
          denominator_value[index] * quotient_value[other] +
          denominator_value[other] * quotient_value[index];
    }
  }

  for (int index = 1; index < table_size; index++) {
    quotient[index] = quotient[index - 1] + quotient_value[index];
  }
  for (u64 block = limit / (split + 1); block > 0; block--) {
    int result_index = table_size - int(block);
    u64 maximum = limit / block;
    u64 root = integer_square_root(maximum);
    quotient[result_index] =
        numerator[result_index] - quotient_value[1] * denominator[result_index] +
        denominator[root] * quotient[root];
    for (u64 left = 2; left <= root; left++) {
      int right_index = denominator.index_of(maximum / left);
      quotient[result_index] -=
          denominator_value[left] * quotient[right_index] +
          quotient_value[left] * denominator[right_index];
    }
  }
  for (int index = 0; index < table_size; index++) {
    quotient[index] *= inverse_constant;
  }
  return quotient;
}

} // namespace dirichlet_prefix_internal

/// @brief Return Q_n={floor(n/i)} in increasing order. Consecutive equal
/// quotients are represented once.
inline std::vector<std::uint64_t> dirichlet_quotients(std::uint64_t limit) {
  dirichlet_prefix_internal::quotient_table<int> layout(limit);
  std::vector<std::uint64_t> result;
  result.reserve(layout.table_size() - 1);
  for (int index = 1; index < layout.table_size(); index++) {
    result.push_back(layout.argument(index));
  }
  return result;
}

/// @brief Given prefix sums of f and g at all increasing values in Q_n,
/// return the corresponding prefix sums of their Dirichlet convolution. Small
/// products are enumerated once around n^(1/3); for a large quotient x, the
/// hyperbola sum groups all equal floor(x/i) values through the Q_n index.
template <class T>
std::vector<T> dirichlet_convolution_prefix_sums(
    std::uint64_t limit, const std::vector<T> &first,
    const std::vector<T> &second) {
  using dirichlet_prefix_internal::convolution;
  using dirichlet_prefix_internal::quotient_table;
  return convolution(quotient_table<T>(limit, first),
                     quotient_table<T>(limit, second))
      .export_values();
}

/// @brief Given prefix sums of f at all increasing values in Q_n, return the
/// prefix sums of its Dirichlet inverse. The recurrence f*g=delta is solved for
/// all small arguments first; each remaining large quotient is then recovered
/// by one grouped divisor-hyperbola equation using those known values.
template <class T>
std::vector<T> dirichlet_inverse_prefix_sums(std::uint64_t limit,
                                             const std::vector<T> &function) {
  using dirichlet_prefix_internal::divide;
  using dirichlet_prefix_internal::quotient_table;
  quotient_table<T> identity(limit);
  for (int index = 1; index < identity.table_size(); index++) {
    identity[index] = T(1);
  }
  return divide(identity, quotient_table<T>(limit, function)).export_values();
}

} // namespace noya