Problem p02851 (C++)

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

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

#define REP(i, n) for (int i = 0; i < (n); ++i)
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define FORR(i, a, b) for (int i = (b)-1; i >= (a); --i)

#define DEBUG(x) cout << #x << ": " << (x) << '\n'
#define DEBUG_VEC(v)                                                           \
  cout << #v << ":";                                                           \
  REP(i, v.size()) cout << ' ' << v[i];                                        \
  cout << '\n'
#define ALL(a) (a).begin(), (a).end()

template <typename T> inline void CHMAX(T &a, const T b) {
  if (a < b)
    a = b;
}
template <typename T> inline void CHMIN(T &a, const T b) {
  if (a > b)
    a = b;
}

constexpr ll MOD = 1000000007ll;
// constexpr ll MOD=998244353ll;
#define FIX(a) ((a) % MOD + MOD) % MOD

const double EPS = 1e-11;
#define EQ0(x) (abs((x)) < EPS)
#define EQ(a, b) (abs((a) - (b)) < EPS)

pll sum[252521];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  // cout<<setprecision(10)<<fixed;
  int n;
  ll k, a[252521];
  cin >> n >> k;
  sum[0] = pll(0, 0);
  REP(i, n) {
    cin >> a[i];
    sum[i + 1] = pll((sum[i].first + a[i]) % k, i + 1);
  }
  REP(i, n) {
    sum[i + 1].first -= i + 1;
    while (sum[i + 1].first < 0) {
      sum[i + 1].first += k;
    }
  }
  sort(sum, sum + n + 1);
  // FOR(i, 0, n+1){
  //   cout<<sum[i].first<<" \n"[i==n];
  // }
  ll ans = 0;
  int r = 0;
  REP(i, n) {
    ll cnt = 0;
    while (r < n + 1 && sum[i].first == sum[r].first &&
           sum[r].second - sum[i].second < k) {
      ++r;
    }
    if (r != i) {
      ans += r - i - 1;
      // cout<<i<<' '<<r<<'\n';
      // DEBUG(ans);
    }
  }
  cout << ans << '\n';
  return 0;
}


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

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

#define REP(i, n) for (int i = 0; i < (n); ++i)
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define FORR(i, a, b) for (int i = (b)-1; i >= (a); --i)

#define DEBUG(x) cout << #x << ": " << (x) << '\n'
#define DEBUG_VEC(v)                                                           \
  cout << #v << ":";                                                           \
  REP(i, v.size()) cout << ' ' << v[i];                                        \
  cout << '\n'
#define ALL(a) (a).begin(), (a).end()

template <typename T> inline void CHMAX(T &a, const T b) {
  if (a < b)
    a = b;
}
template <typename T> inline void CHMIN(T &a, const T b) {
  if (a > b)
    a = b;
}

constexpr ll MOD = 1000000007ll;
// constexpr ll MOD=998244353ll;
#define FIX(a) ((a) % MOD + MOD) % MOD

const double EPS = 1e-11;
#define EQ0(x) (abs((x)) < EPS)
#define EQ(a, b) (abs((a) - (b)) < EPS)

pll sum[252521];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  // cout<<setprecision(10)<<fixed;
  int n;
  ll k, a[252521];
  cin >> n >> k;
  sum[0] = pll(0, 0);
  REP(i, n) {
    cin >> a[i];
    sum[i + 1] = pll((sum[i].first + a[i]) % k, i + 1);
  }
  REP(i, n) {
    sum[i + 1].first -= (i + 1) % k;
    while (sum[i + 1].first < 0) {
      sum[i + 1].first += k;
    }
  }
  sort(sum, sum + n + 1);
  // FOR(i, 0, n+1){
  //   cout<<sum[i].first<<" \n"[i==n];
  // }
  ll ans = 0;
  int r = 0;
  REP(i, n) {
    ll cnt = 0;
    while (r < n + 1 && sum[i].first == sum[r].first &&
           sum[r].second - sum[i].second < k) {
      ++r;
    }
    if (r != i) {
      ans += r - i - 1;
      // cout<<i<<' '<<r<<'\n';
      // DEBUG(ans);
    }
  }
  cout << ans << '\n';
  return 0;
}