Problem p02840 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

#define rep(i, n) for (ll i = 0; i < (n); i++)
#define rep2(i, a, b) for (ll i = (a); i < (b); i++)
#define chmax(x, v)                                                            \
  do {                                                                         \
    x = max(x, v);                                                             \
  } while (0)
typedef uint64_t ull;
typedef int64_t ll;
typedef std::pair<ll, ll> PLL;

using namespace std;

ll N, X, D;

// k個の数字を選んだ時に、合計値として有り得る値の範囲
// 0 <= k <= N
PLL range(ll k) {
  PLL res(k * X, k * X);
  for (ll i = 0; i <= k - 1; i++)
    res.first += i * D;
  for (ll i = N - k; i <= N - 1; i++)
    res.second += i * D;

  return res;
}

map<ll, vector<PLL>> mp;

signed main() {
  cin >> N >> X >> D;

  if (D == 0) {
    if (X == 0) {
      cout << 1 << endl; // S=0のみ
      return 0;
    } else {
      cout << N + 1 << endl; // S=0,X,2X,..,NX
      return 0;
    }
  }

  if (D < 0) {
    X *= -1;
    D *= -1;
  }

  // D > 0
  rep(k, N + 1) {
    ll m = (k * X) % D; // 余り
    m = (m + 3 * D) % D;
    auto r = range(k);
    // printf("k=%d min=%d max=%d m=%d mod_mid=%d mod_max=%d\n",
    //     k, r.first, r.second, m,
    //     (r.first - m)/D, (r.second - m)/D);
    mp[m].emplace_back((r.first - m) / D, (r.second - m) / D);
  }

  // swap(mp[0][0],mp[0][1]);
  // for (auto p: mp[0]) {
  //   printf("mod_min=%d mod_max=%d\n",
  //       p.first, p.second);
  // }

  ll ans = 0;
  for (auto e : mp) {
    ll m = e.first;
    auto ranges = mp[m];
    sort(begin(ranges), end(ranges));
    rep(i, ranges.size()) {
      // printf("mod_min=%d mod_max=%d\n", ranges[i].first, ranges[i].second);
    }

    vector<PLL> merged_ranges;
    merged_ranges.push_back(ranges[0]);
    for (ll i = 1; i < ranges.size(); i++) {
      auto r = ranges[i];
      auto last_r = merged_ranges[merged_ranges.size() - 1];
      if (r.first <= last_r.second + 1) {
        chmax(merged_ranges[merged_ranges.size() - 1].second, r.second);
      } else {
        merged_ranges.push_back(r);
      }
    }

    for (auto mr : merged_ranges) {
      // printf("mr_min=%d mr_max=%d\n", mr.first, mr.second);
      ans += mr.second - mr.first + 1;
    }
  }

  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>

#define rep(i, n) for (ll i = 0; i < (n); i++)
#define rep2(i, a, b) for (ll i = (a); i < (b); i++)
#define chmax(x, v)                                                            \
  do {                                                                         \
    x = max(x, v);                                                             \
  } while (0)
typedef uint64_t ull;
typedef int64_t ll;
typedef std::pair<ll, ll> PLL;

using namespace std;

ll N, X, D;

// k個の数字を選んだ時に、合計値として有り得る値の範囲
// 0 <= k <= N
PLL range(ll k) {
  PLL res(k * X, k * X);

  if (k == N) {
    res.first += ((N - 1) * N / 2) * D;
    res.second += ((N - 1) * N / 2) * D;
  } else {
    res.first += ((k - 1) * k / 2) * D;
    res.second += ((2 * N - k - 1) * k / 2) * D;
  }
  // for (ll i=0; i<=k-1; i++)
  //   res.first += i * D;
  // for (ll i=N-k; i<=N-1; i++)
  //   res.second += i * D;

  return res;
}

map<ll, vector<PLL>> mp;

signed main() {
  cin >> N >> X >> D;

  if (D == 0) {
    if (X == 0) {
      cout << 1 << endl; // S=0のみ
      return 0;
    } else {
      cout << N + 1 << endl; // S=0,X,2X,..,NX
      return 0;
    }
  }

  if (D < 0) {
    X *= -1;
    D *= -1;
  }

  // D > 0
  rep(k, N + 1) {
    ll m = (k * X) % D; // 余り
    m = (m + 3 * D) % D;
    auto r = range(k);
    // printf("k=%d min=%d max=%d m=%d mod_mid=%d mod_max=%d\n",
    //     k, r.first, r.second, m,
    //     (r.first - m)/D, (r.second - m)/D);
    mp[m].emplace_back((r.first - m) / D, (r.second - m) / D);
  }

  // swap(mp[0][0],mp[0][1]);
  // for (auto p: mp[0]) {
  //   printf("mod_min=%d mod_max=%d\n",
  //       p.first, p.second);
  // }

  ll ans = 0;
  for (auto e : mp) {
    ll m = e.first;
    auto ranges = mp[m];
    sort(begin(ranges), end(ranges));
    rep(i, ranges.size()) {
      // printf("mod_min=%d mod_max=%d\n", ranges[i].first, ranges[i].second);
    }

    vector<PLL> merged_ranges;
    merged_ranges.push_back(ranges[0]);
    for (ll i = 1; i < ranges.size(); i++) {
      auto r = ranges[i];
      auto last_r = merged_ranges[merged_ranges.size() - 1];
      if (r.first <= last_r.second + 1) {
        chmax(merged_ranges[merged_ranges.size() - 1].second, r.second);
      } else {
        merged_ranges.push_back(r);
      }
    }

    for (auto mr : merged_ranges) {
      // printf("mr_min=%d mr_max=%d\n", mr.first, mr.second);
      ans += mr.second - mr.first + 1;
    }
  }

  cout << ans << endl;
  return 0;
}