Problem p02919 (C++)

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

using namespace std;

using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using vi = vector<int>;
using vl = vector<ll>;

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

#define CHMIN(a, b) (a) = min((a), (b))
#define CHMAX(a, b) (a) = max((a), (b))
#define ALL(v) (v).begin(), (v).end()

#define DEBUG(x) cerr << #x << ": " << (x) << endl
#define DEBUG_VEC(v)                                                           \
  cerr << #v << ": ";                                                          \
  REP(__i, (v).size()) cerr << ((v)[__i]) << ", ";                             \
  cerr << endl

const ll MOD = 1000000007ll;

int n;
int p[125252];
int rev[125252];

const int N = 1 << 17;
pii dat[2 * N];

pii merge(pii a, pii b) {
  int x[4] = {a.first, a.second, b.first, b.second};
  sort(x, x + 4);
  return pii(x[3], x[2]);
}
pii query(int l, int r, int a, int b, int k) {
  if (r <= a || b <= l)
    return pii(-1, -1);
  if (l <= a && b <= r)
    return dat[k];
  int m = (a + b) / 2;
  pii x = query(l, r, a, m, 2 * k + 1);
  pii y = query(l, r, m, b, 2 * k + 2);
  return merge(x, y);
}

int main() {
  scanf("%d", &n);
  REP(i, n) scanf("%d", p + i);
  REP(i, n) rev[p[i]] = i;
  REP(i, n) dat[i + N - 1] = pii(p[i], -1);
  FORR(i, 0, N - 1) { dat[i] = merge(dat[2 * i + 1], dat[2 * i + 2]); }
  ll ans = 0;
  REP(i, n) {
    int val = p[i];
    if (val == n)
      continue;
    if (query(0, i + 1, 0, N, 0).first == val) {
      // left has no max
      int ok = i + 2, ng = n + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        pii P = query(i, r, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = r;
      }
      int r2 = ok;
      ok = r2;
      ng = i + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        (query(i, r, 0, N, 0).second == val ? ok : ng) = r;
      }
      int r1 = ok;
      r2--;
      r1--;
      ans += 1ll * val * (i + 1) * (r2 - r1 + 1);
    } else if (query(i, n, 0, N, 0).first == val) {
      // right has no max
      int ok = i - 1, ng = -1;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        pii P = query(l, i + 1, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = l;
      }
      int l2 = ok;
      ok = l2;
      ng = i;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        (query(l, i + 1, 0, N, 0).second == val ? ok : ng) = l;
      }
      int l1 = ok;
      ans += 1ll * val * (n - i) * (l1 - l2 + 1);
    } else {
      int ok = i + 2, ng = n + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        pii P = query(i, r, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = r;
      }
      int r2 = ok;
      ok = r2;
      ng = i + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        (query(i, r, 0, N, 0).second == val ? ok : ng) = r;
      }
      int r1 = ok;
      r2--;
      r1--;
      ok = i - 1, ng = -1;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        pii P = query(l, i + 1, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = l;
      }
      int l2 = ok;
      ok = l2;
      ng = i;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        (query(l, i + 1, 0, N, 0).second == val ? ok : ng) = l;
      }
      int l1 = ok;
      ans += 1ll * val * (r2 - r1 + 1) * (i - l1);
      ans += 1ll * val * (l1 - l2 + 1) * (r1 - i);
    }
  }
  printf("%lld\n", ans);
  return 0;
}


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

using namespace std;

using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using vi = vector<int>;
using vl = vector<ll>;

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

#define CHMIN(a, b) (a) = min((a), (b))
#define CHMAX(a, b) (a) = max((a), (b))
#define ALL(v) (v).begin(), (v).end()

#define DEBUG(x) cerr << #x << ": " << (x) << endl
#define DEBUG_VEC(v)                                                           \
  cerr << #v << ": ";                                                          \
  REP(__i, (v).size()) cerr << ((v)[__i]) << ", ";                             \
  cerr << endl

const ll MOD = 1000000007ll;

int n;
int p[125252];
int rev[125252];

const int N = 1 << 17;
pii dat[2 * N];

pii merge(pii a, pii b) {
  if (a.first > b.first) {
    return pii(a.first, max(a.second, b.first));
  } else {
    return pii(b.first, max(a.first, b.second));
  }
}
pii query(int l, int r, int a, int b, int k) {
  if (r <= a || b <= l)
    return pii(-1, -1);
  if (l <= a && b <= r)
    return dat[k];
  int m = (a + b) / 2;
  pii x = query(l, r, a, m, 2 * k + 1);
  pii y = query(l, r, m, b, 2 * k + 2);
  return merge(x, y);
}

int main() {
  scanf("%d", &n);
  REP(i, n) scanf("%d", p + i);
  REP(i, n) rev[p[i]] = i;
  REP(i, n) dat[i + N - 1] = pii(p[i], -1);
  FORR(i, 0, N - 1) { dat[i] = merge(dat[2 * i + 1], dat[2 * i + 2]); }
  ll ans = 0;
  REP(i, n) {
    int val = p[i];
    if (val == n)
      continue;
    if (query(0, i + 1, 0, N, 0).first == val) {
      // left has no max
      int ok = i + 2, ng = n + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        pii P = query(i, r, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = r;
      }
      int r2 = ok;
      ok = r2;
      ng = i + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        (query(i, r, 0, N, 0).second == val ? ok : ng) = r;
      }
      int r1 = ok;
      r2--;
      r1--;
      ans += 1ll * val * (i + 1) * (r2 - r1 + 1);
    } else if (query(i, n, 0, N, 0).first == val) {
      // right has no max
      int ok = i - 1, ng = -1;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        pii P = query(l, i + 1, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = l;
      }
      int l2 = ok;
      ok = l2;
      ng = i;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        (query(l, i + 1, 0, N, 0).second == val ? ok : ng) = l;
      }
      int l1 = ok;
      ans += 1ll * val * (n - i) * (l1 - l2 + 1);
    } else {
      int ok = i + 2, ng = n + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        pii P = query(i, r, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = r;
      }
      int r2 = ok;
      ok = r2;
      ng = i + 1;
      while (abs(ok - ng) > 1) {
        int r = (ok + ng) / 2;
        (query(i, r, 0, N, 0).second == val ? ok : ng) = r;
      }
      int r1 = ok;
      r2--;
      r1--;
      ok = i - 1, ng = -1;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        pii P = query(l, i + 1, 0, N, 0);
        (P.first == val || P.second == val ? ok : ng) = l;
      }
      int l2 = ok;
      ok = l2;
      ng = i;
      while (abs(ok - ng) > 1) {
        int l = (ok + ng) / 2;
        (query(l, i + 1, 0, N, 0).second == val ? ok : ng) = l;
      }
      int l1 = ok;
      ans += 1ll * val * (r2 - r1 + 1) * (i - l1);
      ans += 1ll * val * (l1 - l2 + 1) * (r1 - i);
    }
  }
  printf("%lld\n", ans);
  return 0;
}