Problem p01913 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
#define repl(i, a, b) for (int i = (int)(a); i < (int)(b); i++)
#define rep(i, n) repl(i, 0, n)
#define mp(a, b) make_pair((a), (b))
#define pb(a) push_back((a))
#define all(x) (x).begin(), (x).end()
#define uniq(x) sort(all(x)), (x).erase(unique(all(x)), end(x))
#define fi first
#define se second
#define dbg(...) _dbg(#__VA_ARGS__ ",", __VA_ARGS__)
void _dbg(string) { cout << endl; }
template <class H, class... T> void _dbg(string s, H h, T... t) {
  int l = s.find(',');
  cout << s.substr(0, l) << " = " << h << ", ";
  _dbg(s.substr(l + 1), t...);
}
template <class T, class U>
ostream &operator<<(ostream &o, const pair<T, U> &p) {
  o << "(" << p.fi << "," << p.se << ")";
  return o;
}
template <class T> ostream &operator<<(ostream &o, const vector<T> &v) {
  o << "[";
  for (T t : v) {
    o << t << ",";
  }
  o << "]";
  return o;
}

#define INF 2147483600

class UnionFind {
public:
  vector<int> par,
      rank; // parent(negative := its root and abs-value is its size), depth
  UnionFind(int sz) : par(sz, -1), rank(sz, 0) {}
  int find(int x) {
    if (par[x] < 0)
      return x;
    else
      return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;                // already belong to same tree
    if (rank[x] < rank[y]) { // y becomes parent node
      par[y] += par[x];
      par[x] = y;
    } else { // x becomes parent node
      par[x] += par[y];
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int size(int x) { return -par[find(x)]; }
};

int main() {
  int n, m, T;
  cin >> n >> m >> T;
  vector<vector<pair<int, int>>> vec(n);
  vector<pair<int, pair<int, int>>> edge(m);
  rep(i, m) {
    int a, b, t;
    cin >> a >> b >> t;
    a--;
    b--;
    vec[a].pb(mp(b, t));
    vec[b].pb(mp(a, t));
    edge[i] = mp(t, mp(a, b));
  }

  // dijkstra??§??°?????????????????????
  vector<int> d(n, INF);
  using P = pair<int, int>;
  priority_queue<P, vector<P>, greater<P>> pq;
  pq.push(mp(0, 0));
  d[0] = 0;
  while (!pq.empty()) {
    auto p = pq.top();
    pq.pop();
    int v = p.second;
    int dis = p.first;
    if (d[v] < dis)
      continue;
    for (auto to : vec[v]) {
      if (to.se <= dis)
        continue;
      if (d[to.first] <= dis + 1)
        continue;
      d[to.first] = dis + 1;
      pq.push(mp(dis + 1, to.first));
    }
  }

  // d[i]<INF ????????°?????????

  sort(all(edge));
  reverse(all(edge));
  UnionFind uf(n);

  vector<long> maxval(n, 0);
  vector<int> lastvisit(n, T);
  rep(i, n) if (d[i] < INF) maxval[i] = 1;

  auto getlazy = [&](int x, int t) {
    int node = uf.find(x);
    if (maxval[node] == 0)
      return 0L;
    else if (lastvisit[node] == t)
      return maxval[node];
    else {
      long add = (long)(lastvisit[node] - t) * uf.size(x);
      return maxval[node] + add;
    }
  };

  rep(i, m) {
    int u = edge[i].se.fi, v = edge[i].se.se;
    int t = edge[i].fi;
    if (uf.same(u, v))
      continue;
    else if (t == T) {
      bool visitable = true;
      if (maxval[uf.find(u)] == 0 && maxval[uf.find(v)] == 0)
        visitable = false;
      uf.unite(u, v);
      if (visitable)
        maxval[uf.find(u)] = uf.size(u);
    } else {
      long lu = getlazy(u, t), lv = getlazy(v, t);
      if (lu == 0 && lv == 0) {
        uf.unite(u, v);
        continue;
      }
      if (lu > 0)
        lu += uf.size(v);
      if (lv > 0)
        lv += uf.size(u);
      uf.unite(u, v);
      maxval[uf.find(u)] = max<long>(lu, lv);
      lastvisit[uf.find(u)] = t; // dbg(uf.find(u), maxval[uf.find(u)]);
    }
  }

  cout << getlazy(0, 1) << endl;
  assert(getlazy(0, 1) > 0);

  return 0;
}

// WA memo
// ???lu==0 && lv==0
// ?????¨??????unite????????????????????¨??????????????????????????????
// ???t=1??§???unite???????????????????????´???????????????lazy???????????????????????????????°????????????£???????????£???
// ???Dijkstra???????????§=??§???priority_queue????????£????????§?????????MLE??????
// ??¨?????????dijkstra???????????????????????????1???????????????????????????best

Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define repl(i, a, b) for (int i = (int)(a); i < (int)(b); i++)
#define rep(i, n) repl(i, 0, n)
#define mp(a, b) make_pair((a), (b))
#define pb(a) push_back((a))
#define all(x) (x).begin(), (x).end()
#define uniq(x) sort(all(x)), (x).erase(unique(all(x)), end(x))
#define fi first
#define se second
#define dbg(...) _dbg(#__VA_ARGS__ ",", __VA_ARGS__)
void _dbg(string) { cout << endl; }
template <class H, class... T> void _dbg(string s, H h, T... t) {
  int l = s.find(',');
  cout << s.substr(0, l) << " = " << h << ", ";
  _dbg(s.substr(l + 1), t...);
}
template <class T, class U>
ostream &operator<<(ostream &o, const pair<T, U> &p) {
  o << "(" << p.fi << "," << p.se << ")";
  return o;
}
template <class T> ostream &operator<<(ostream &o, const vector<T> &v) {
  o << "[";
  for (T t : v) {
    o << t << ",";
  }
  o << "]";
  return o;
}

#define INF 2147483600

class UnionFind {
public:
  vector<int> par,
      rank; // parent(negative := its root and abs-value is its size), depth
  UnionFind(int sz) : par(sz, -1), rank(sz, 0) {}
  int find(int x) {
    if (par[x] < 0)
      return x;
    else
      return par[x] = find(par[x]);
  }
  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;                // already belong to same tree
    if (rank[x] < rank[y]) { // y becomes parent node
      par[y] += par[x];
      par[x] = y;
    } else { // x becomes parent node
      par[x] += par[y];
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }
  bool same(int x, int y) { return find(x) == find(y); }
  int size(int x) { return -par[find(x)]; }
};

int main() {
  int n, m, T;
  cin >> n >> m >> T;
  vector<vector<pair<int, int>>> vec(n);
  vector<pair<int, pair<int, int>>> edge(m);
  rep(i, m) {
    int a, b, t;
    cin >> a >> b >> t;
    a--;
    b--;
    vec[a].pb(mp(b, t));
    vec[b].pb(mp(a, t));
    edge[i] = mp(t, mp(a, b));
  }

  // dijkstra??§??°?????????????????????
  vector<int> d(n, INF);
  using P = pair<int, int>;
  priority_queue<P, vector<P>, greater<P>> pq;
  pq.push(mp(0, 0));
  d[0] = 0;
  while (!pq.empty()) {
    auto p = pq.top();
    pq.pop();
    int v = p.second;
    int dis = p.first;
    if (d[v] < dis)
      continue;
    for (auto to : vec[v]) {
      if (to.se <= dis)
        continue;
      if (d[to.first] < dis + 1)
        continue;
      d[to.first] = dis + 1;
      pq.push(mp(dis + 1, to.first));
    }
  }

  vector<vector<pair<int, int>>>().swap(vec); // free memory

  // d[i]<INF ????????°?????????

  sort(all(edge));
  reverse(all(edge));
  UnionFind uf(n);

  vector<long> maxval(n, 0);
  vector<int> lastvisit(n, T);
  rep(i, n) if (d[i] < INF) maxval[i] = 1;

  auto getlazy = [&](int x, int t) {
    int node = uf.find(x);
    if (maxval[node] == 0)
      return 0L;
    else if (lastvisit[node] == t)
      return maxval[node];
    else {
      long add = (lastvisit[node] - t) * uf.size(x);
      return maxval[node] + add;
    }
  };

  rep(i, m) {
    int u = edge[i].se.fi, v = edge[i].se.se;
    int t = edge[i].fi;
    if (uf.same(u, v))
      continue;
    else if (t == T) {
      bool visitable = true;
      if (maxval[uf.find(u)] == 0 && maxval[uf.find(v)] == 0)
        visitable = false;
      uf.unite(u, v);
      if (visitable)
        maxval[uf.find(u)] = uf.size(u);
    } else {
      long lu = getlazy(u, t), lv = getlazy(v, t);
      if (lu == 0 && lv == 0) {
        uf.unite(u, v);
        continue;
      }
      lu += uf.size(v);
      lv += uf.size(u);
      uf.unite(u, v);
      maxval[uf.find(u)] = max<long>(lu, lv);
      lastvisit[uf.find(u)] = t; // dbg(uf.find(u), maxval[uf.find(u)]);
    }
  }

  cout << getlazy(0, 1) << endl;

  return 0;
}