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

#include "bits/stdc++.h"
using namespace std;

struct edge {
  int from;
  int to;
  int t;
};

int N, M, T;
void bfs(const vector<vector<edge>> &edges, vector<int> &can_go) {
  queue<pair<int, int>> que;
  que.emplace(0, 0);
  can_go[0] = true;
  while (!que.empty()) {
    auto atop(que.front());
    que.pop();
    const int nexttime = atop.second + 1;
    for (auto e : edges[atop.first]) {
      if (nexttime <= e.t && !can_go[e.to]) {
        can_go[e.to] = true;
        que.emplace(e.to, nexttime);
      }
    }
  }
}

struct UnionFind {
  vector<int> data;
  UnionFind(int size) : data(size, -1) {}
  bool unionSet(int x, int y) {
    x = root(x);
    y = root(y);
    if (x != y) {
      if (data[y] < data[x])
        swap(x, y);
      data[x] += data[y];
      data[y] = x;
    }
    return x != y;
  }
  bool findSet(int x, int y) { return root(x) == root(y); }
  int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]); }
  int size(int x) { return -data[root(x)]; }
};

int main() {
  cin >> N >> M >> T;
  vector<vector<edge>> edges(N);
  vector<edge> all_edges;
  for (int i = 0; i < M; ++i) {
    int a, b, t;
    cin >> a >> b >> t;
    a--;
    b--;
    edges[a].push_back(edge{a, b, t});
    edges[b].push_back(edge{b, a, t});
    all_edges.push_back(edge{a, b, t});
  }
  sort(all_edges.begin(), all_edges.end(),
       [](const edge &l, const edge &r) { return l.t > r.t; });
  vector<long long int> scores(N);
  vector<long long int> latest_connect(N, T);
  UnionFind uf(N);
  vector<int> can_go(N);
  bfs(edges, can_go);
  for (auto e : all_edges) {
    int a(e.from);
    int b(e.to);
    const int root_a = uf.root(a);
    const int root_b = uf.root(b);
    if (root_a == root_b)
      continue;
    else {
      long long int new_score = 0;
      if (can_go[root_a]) {
        new_score =
            max(new_score, scores[root_a] + (latest_connect[root_a] - e.t) *
                                                uf.size(root_a));
      }
      if (can_go[root_b]) {
        new_score =
            max(new_score, scores[root_b] + (latest_connect[root_b] - e.t) *
                                                uf.size(root_b));
      }
      uf.unionSet(root_a, root_b);
      const int root_n = uf.root(root_a);
      if (can_go[root_a]) {
        can_go[root_n] = true;
      }
      if (can_go[root_b]) {
        can_go[root_n] = true;
      }
      latest_connect[root_n] = e.t;
      scores[root_n] = new_score;
    }
  }
  const int aroot = uf.root(0);
  long long int ans = scores[aroot] + latest_connect[aroot] * N;
  cout << ans << endl;
  return 0;
}

Error: MLE

Fixed submission:

#include "bits/stdc++.h"
using namespace std;

struct edge {
  int from;
  int to;
  int t;
};

int N, M, T;
void bfs(vector<vector<edge>> &edges, vector<int> &can_go) {
  queue<pair<int, int>> que;
  que.emplace(0, 0);
  can_go[0] = true;
  while (!que.empty()) {
    auto atop(que.front());
    que.pop();
    const int nexttime = atop.second + 1;
    for (auto e : edges[atop.first]) {
      if (nexttime <= e.t) {
        can_go[e.to] = true;
        que.emplace(e.to, nexttime);
      }
    }
  }
}

struct UnionFind {
  vector<int> data;
  UnionFind(int size) : data(size, -1) {}
  bool unionSet(int x, int y) {
    x = root(x);
    y = root(y);
    if (x != y) {
      if (data[y] < data[x])
        swap(x, y);
      data[x] += data[y];
      data[y] = x;
    }
    return x != y;
  }
  bool findSet(int x, int y) { return root(x) == root(y); }
  int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]); }
  int size(int x) { return -data[root(x)]; }
};

int main() {
  cin >> N >> M >> T;
  vector<vector<edge>> edges(N);
  vector<edge> all_edges;
  for (int i = 0; i < M; ++i) {
    int a, b, t;
    cin >> a >> b >> t;
    a--;
    b--;
    edges[a].push_back(edge{a, b, t});
    edges[b].push_back(edge{b, a, t});
    all_edges.push_back(edge{a, b, t});
  }
  sort(all_edges.begin(), all_edges.end(),
       [](const edge &l, const edge &r) { return l.t > r.t; });
  vector<long long int> scores(N);
  vector<long long int> latest_connect(N, T);
  UnionFind uf(N);
  vector<int> can_go(N);
  bfs(edges, can_go);
  for (auto e : all_edges) {
    int a(e.from);
    int b(e.to);
    const int root_a = uf.root(a);
    const int root_b = uf.root(b);
    if (root_a == root_b)
      continue;
    else {
      long long int new_score = 0;
      if (can_go[root_a]) {
        new_score =
            max(new_score, scores[root_a] + (latest_connect[root_a] - e.t) *
                                                uf.size(root_a));
      }
      if (can_go[root_b]) {
        new_score =
            max(new_score, scores[root_b] + (latest_connect[root_b] - e.t) *
                                                uf.size(root_b));
      }
      uf.unionSet(root_a, root_b);
      const int root_n = uf.root(root_a);
      if (can_go[root_a]) {
        can_go[root_n] = true;
      }
      if (can_go[root_b]) {
        can_go[root_n] = true;
      }
      latest_connect[root_n] = e.t;
      scores[root_n] = new_score;
    }
  }
  const int aroot = uf.root(0);
  long long int ans = scores[aroot] + latest_connect[aroot] * N;
  cout << ans << endl;
  return 0;
}