tests/small_range_lower_bound.test.cpp

• View this file on GitHub
• Last update: 2025-01-23 04:00:37+00:00
• Problem: https://atcoder.jp/contests/abc356/tasks/abc356_e

Depends on

• std::lower_bound in $\left\langle O(N + \max(A_i) - \min(A_i)), O(1) \right\rangle$ time (tools/small_range_lower_bound.hpp)

Code

// competitive-verifier: PROBLEM https://atcoder.jp/contests/abc356/tasks/abc356_e
// competitive-verifier: IGNORE

#include <iostream>
#include <vector>
#include <algorithm>
#include "tools/small_range_lower_bound.hpp"

using ll = long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N;
  std::cin >> N;
  std::vector<int> A(N);
  for (auto&& A_i : A) std::cin >> A_i;

  std::sort(A.begin(), A.end());
  tools::small_range_lower_bound<int> lower_bound(A.begin(), A.end());

  ll answer = 0;
  for (int l = 0, r = 0; l < N; l = r) {
    for (; r < N && A[l] == A[r]; ++r);
    for (int q = 1, max_q = A.back() / A[l]; q <= max_q; ++q) {
      answer += static_cast<ll>(lower_bound.query(A[l] * (q + 1)) - lower_bound.query(A[l] * q)) * q * (r - l);
    }
    answer -= static_cast<ll>(r - l) * (r - l + 1) / 2;
  }

  std::cout << answer << '\n';
  return 0;
}

#line 1 "tests/small_range_lower_bound.test.cpp"
// competitive-verifier: PROBLEM https://atcoder.jp/contests/abc356/tasks/abc356_e
// competitive-verifier: IGNORE

#include <iostream>
#include <vector>
#include <algorithm>
#line 1 "tools/small_range_lower_bound.hpp"



#line 5 "tools/small_range_lower_bound.hpp"
#include <type_traits>
#include <iterator>
#include <cstddef>
#include <limits>
#line 10 "tools/small_range_lower_bound.hpp"
#include <numeric>

namespace tools {
  template <typename T>
  class small_range_lower_bound {
    T m_min;
    ::std::vector<int> m_res;

    template <typename ForwardIterator, ::std::enable_if_t<::std::is_base_of_v<::std::forward_iterator_tag, typename ::std::iterator_traits<ForwardIterator>::iterator_category>, ::std::nullptr_t> = nullptr>
    void init(const ForwardIterator begin, const ForwardIterator end) {
      if (begin == end) {
        this->m_min = ::std::numeric_limits<T>::max();
        this->m_res.assign({0});
      } else {
        const auto [minit, maxit] = ::std::minmax_element(begin, end);
        this->m_min = *minit;
        this->m_res.assign(*maxit - *minit + 2, 0);
        for (auto it = begin; it != end; ++it) {
          ++this->m_res[*it - *minit + 1];
        }
        ::std::partial_sum(this->m_res.begin(), this->m_res.end(), this->m_res.begin());
      }
    }

  public:
    small_range_lower_bound() = default;
    template <typename ForwardIterator, ::std::enable_if_t<::std::is_base_of_v<::std::forward_iterator_tag, typename ::std::iterator_traits<ForwardIterator>::iterator_category>, ::std::nullptr_t> = nullptr>
    small_range_lower_bound(const ForwardIterator begin, const ForwardIterator end) {
      this->init(begin, end);
    }
    template <typename InputIterator, ::std::enable_if_t<!::std::is_base_of_v<::std::forward_iterator_tag, typename ::std::iterator_traits<InputIterator>::iterator_category>, ::std::nullptr_t> = nullptr>
    small_range_lower_bound(const InputIterator begin, const InputIterator end) {
      ::std::vector<T> v(begin, end);
      this->init(v.begin(), v.end());
    }

    T query(const T x) {
      if (x < this->m_min) return 0;
      return this->m_res[::std::min<int>(x - this->m_min, this->m_res.size() - 1)];
    }
  };
}


#line 8 "tests/small_range_lower_bound.test.cpp"

using ll = long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N;
  std::cin >> N;
  std::vector<int> A(N);
  for (auto&& A_i : A) std::cin >> A_i;

  std::sort(A.begin(), A.end());
  tools::small_range_lower_bound<int> lower_bound(A.begin(), A.end());

  ll answer = 0;
  for (int l = 0, r = 0; l < N; l = r) {
    for (; r < N && A[l] == A[r]; ++r);
    for (int q = 1, max_q = A.back() / A[l]; q <= max_q; ++q) {
      answer += static_cast<ll>(lower_bound.query(A[l] * (q + 1)) - lower_bound.query(A[l] * q)) * q * (r - l);
    }
    answer -= static_cast<ll>(r - l) * (r - l + 1) / 2;
  }

  std::cout << answer << '\n';
  return 0;
}

Back to top page