Codeforces/1829/g.cpp

#include <algorithm>
#include <cassert>
#include <cctype>
#include <cstdint>
#include <functional>
#include <iostream>
#include <map>
#include <optional>
#include <queue>
#include <set>
#include <sstream>
#include <string>
#include <vector>

template <typename T, typename U>
std::ostream &operator<<(std::ostream &s, std::pair<T, U> const &p) {
  return s << p.first << " " << p.second;
}

namespace helpers {

using namespace std;

namespace math {

uint32_t pow(uint32_t base, uint32_t exp) {
  if (exp == 0) return 1;
  uint32_t half = pow(base, exp / 2);
  if (exp % 2 == 0) return half * half;
  return half * half * base;
}

}  // namespace math

namespace input {

template <typename T>
std::vector<T> load_vector(std::size_t size) {
  std::vector<T> result{};

  for (auto i = 0u; i < size; ++i) {
    T x;
    std::cin >> x;
    result.push_back(std::move(x));
  }

  return result;
}

}  // namespace input

namespace output {

template <typename T>
inline void answer(const T &ans) {
  cout << ans << "\n";
}

inline void yes() { cout << "YES\n"; }
inline void no() { cout << "NO\n"; }

inline void yesno(bool ans) {
  if (ans) {
    yes();
  } else {
    no();
  }
}

}  // namespace output

using namespace math;
using namespace input;
using namespace output;

#define LOOP(n) for (auto i = 0; i < n; ++i)

}  // namespace helpers

// for ‹N› test cases, uncomment for single test case
// #define SINGLE

namespace solution {

using namespace std;
using namespace helpers;

constexpr array<int32_t, 2023> ROWS = []() {
  array<int32_t, 2023> starts;

  starts[0] = 1;
  for (auto y = 1; y < 2023; ++y) {
    starts[y] = starts[y - 1] + y;
  }

  return starts;
}();

int32_t find_row(int n) {
  return distance(ROWS.begin(), upper_bound(ROWS.begin(), ROWS.end(), n)) - 1;
}

int32_t get_n(int y, int x) { return ROWS[y] + x; }

uint64_t sum_of_n_squared(uint64_t n) { return n * (n + 1) * (2 * n + 1) / 6; }

uint64_t add_row(int32_t lower, int32_t upper) {
  return sum_of_n_squared(upper) - sum_of_n_squared(lower - 1);
}

uint64_t find_sum(int32_t n) {
  auto y = find_row(n);
  auto x_u = n - ROWS[y];
  auto x_l = x_u;

  uint64_t sum = 0;
  for (; y >= 0; --y) {
    auto lower = get_n(y, x_l);
    auto upper = get_n(y, x_u);

    sum += add_row(lower, upper);

    x_l = max(0, x_l - 1);
    x_u = min(x_u, y - 1);
  }

  return sum;
}

void solve() {
  int n;
  cin >> n;
  answer(find_sum(n));
}

}  // namespace solution

using namespace solution;

#ifdef TEST

#include "../.common/cpp/catch_amalgamated.hpp"

TEST_CASE("find row") {
  CHECK(ROWS[find_row(1)] == 1);

  CHECK(ROWS[find_row(2)] == 2);
  CHECK(ROWS[find_row(3)] == 2);

  for (int i = 11; i <= 15; ++i) {
    CHECK(ROWS[find_row(i)] == 11);
  }
}

TEST_CASE("examples") {
  CHECK(find_sum(9) == 156);
  CHECK(find_sum(1) == 1);
  CHECK(find_sum(2) == 5);
  CHECK(find_sum(3) == 10);
  CHECK(find_sum(4) == 21);
  CHECK(find_sum(5) == 39);
  CHECK(find_sum(6) == 46);
  CHECK(find_sum(10) == 146);
  CHECK(find_sum(1434) == 63145186);

  LOOP(1000) { CHECK(find_sum(1000000) == 58116199242129511); }
  LOOP(100000) { find_sum(900000 + i); }
}

#else

int main(void) {

#ifdef SINGLE

  solution::solve();

#else

  // for multiple test cases
  int N;
  std::cin >> N >> std::ws;

  for (auto i = 0; i < N; ++i) {
    solution::solve();
  }

#endif

  return 0;
}

#endif