Problem p01827 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<ll, ll> P;

#define each(i, a) for (auto &&i : a)
#define FOR(i, a, b) for (ll i = (a), __last_##i = (b); i < __last_##i; i++)
#define RFOR(i, a, b) for (ll i = (b)-1, __last_##i = (a); i >= __last_##i; i--)
#define REP(i, n) FOR(i, 0, n)
#define RREP(i, n) RFOR(i, 0, n)
#define __GET_MACRO3(_1, _2, _3, NAME, ...) NAME
#define rep(...) __GET_MACRO3(__VA_ARGS__, FOR, REP)(__VA_ARGS__)
#define rrep(...) __GET_MACRO3(__VA_ARGS__, RFOR, RREP)(__VA_ARGS__)
#define pb push_back
#define all(a) (a).begin(), (a).end()
#define chmin(x, v) x = min(x, v)
#define chmax(x, v) x = max(x, v)

const ll linf = 1e18;
const int inf = 1e9;
const double eps = 1e-12;
const double pi = acos(-1);

template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  each(x, vec) is >> x;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  rep(i, vec.size()) {
    if (i)
      os << " ";
    os << vec[i];
  }
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const vector<vector<T>> &vec) {
  rep(i, vec.size()) {
    if (i)
      os << endl;
    os << vec[i];
  }
  return os;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  ll n;
  cin >> n;
  vector<ll> c(n);
  cin >> c;
  ll x;
  cin >> x;
  vector<vector<ll>> G(n);
  rep(i, x) {
    ll a, b;
    cin >> a >> b;
    --a, --b;
    G[a].pb(b);
    G[b].pb(a);
  }
  map<ll, vector<ll>> m;
  rep(i, n) m[c[i]].pb(i);
  vector<ll> ans(n, 0);
  vector<ll> M(n, 0);
  vector<queue<ll>> Q(n);
  vector<bool> used(n, false);
  each(p, m) {
    // vector<ll> nans = ans;
    // vector<ll> nM = M;
    vector<P> a;
    each(v, p.second) {
      ll val = 1;
      chmax(val, M[v]);
      each(to, G[v]) {
        chmax(val, M[to] + 1);
        Q[to].push(v);
        // chmax(val, ans[to]+1);
      }
      ans[v] = val;
      a.pb({val, v});
    }
    sort(all(a), greater<P>());
    each(pp, a) {
      ll val = pp.first;
      queue<ll> q;
      q.push(pp.second);
      while (!q.empty()) {
        ll v = q.front();
        q.pop();
        each(to, G[v]) {
          if (c[to] == c[v]) {
            if (used[to])
              continue;
            used[to] = true;
            chmax(ans[to], val);
            q.push(to);
          } else {
            while (!Q[to].empty()) {
              ll nv = Q[to].front();
              Q[to].pop();
              if (used[nv])
                continue;
              used[nv] = true;
              chmax(ans[nv], val);
              q.push(nv);
            }
          }
        }
      }
    }
    // each(v, p.second) {
    //     ll val = 1;
    //     chmax(val, nM[v]);
    //     each(to, G[v]) {
    //         chmax(val, M[to]+1);
    //         chmax(val, ans[to]+1);
    //     }
    //     nans[v] = val;
    //     chmax(nM[v], val);
    //     each(to, G[v]) {
    //         chmax(nM[to], val);
    //     }
    // }
    each(v, p.second) {
      chmax(M[v], ans[v]);
      each(to, G[v]) { chmax(M[to], ans[v]); }
    }
    // M = nM;
    // ans = nans;
  }
  // cout << ans << endl;
  ll sum = 0;
  rep(i, n) sum += ans[i];
  cout << sum << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<ll, ll> P;

#define each(i, a) for (auto &&i : a)
#define FOR(i, a, b) for (ll i = (a), __last_##i = (b); i < __last_##i; i++)
#define RFOR(i, a, b) for (ll i = (b)-1, __last_##i = (a); i >= __last_##i; i--)
#define REP(i, n) FOR(i, 0, n)
#define RREP(i, n) RFOR(i, 0, n)
#define __GET_MACRO3(_1, _2, _3, NAME, ...) NAME
#define rep(...) __GET_MACRO3(__VA_ARGS__, FOR, REP)(__VA_ARGS__)
#define rrep(...) __GET_MACRO3(__VA_ARGS__, RFOR, RREP)(__VA_ARGS__)
#define pb push_back
#define all(a) (a).begin(), (a).end()
#define chmin(x, v) x = min(x, v)
#define chmax(x, v) x = max(x, v)

const ll linf = 1e18;
const int inf = 1e9;
const double eps = 1e-12;
const double pi = acos(-1);

template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  each(x, vec) is >> x;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  rep(i, vec.size()) {
    if (i)
      os << " ";
    os << vec[i];
  }
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const vector<vector<T>> &vec) {
  rep(i, vec.size()) {
    if (i)
      os << endl;
    os << vec[i];
  }
  return os;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  ll n;
  cin >> n;
  vector<ll> c(n);
  cin >> c;
  ll x;
  cin >> x;
  vector<vector<ll>> G(n);
  rep(i, x) {
    ll a, b;
    cin >> a >> b;
    --a, --b;
    G[a].pb(b);
    G[b].pb(a);
  }
  map<ll, vector<ll>> m;
  rep(i, n) m[c[i]].pb(i);
  vector<ll> ans(n, 0);
  vector<ll> M(n, 0);
  each(p, m) {
    vector<ll> nans = ans;
    vector<ll> nM = M;
    each(v, p.second) {
      ll val = 1;
      chmax(val, nM[v]);
      each(to, G[v]) {
        chmax(val, M[to] + 1);
        chmax(val, ans[to] + 1);
      }
      nans[v] = val;
      chmax(nM[v], val);
      each(to, G[v]) { chmax(nM[to], val); }
    }
    M = nM;
    ans = nans;
  }
  ll sum = 0;
  rep(i, n) sum += ans[i];
  cout << sum << endl;
}