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

#include <bits/stdc++.h>
#define range(i, a, b) for (int i = (a); i < (b); i++)
#define rep(i, b) for (int i = 0; i < (b); i++)
#define all(a) (a).begin(), (a).end()
#define show(x) cerr << #x << " = " << (x) << endl;
using namespace std;

template <typename T> ostream &operator<<(ostream &os, vector<T> &v) {
  rep(i, v.size()) { os << v[i] << (i == v.size() - 1 ? "" : " "); }
  return os;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &x : v) {
    is >> x;
  }
  return is;
}

template <typename T> struct sumGroup {
  using value_type = T;
  constexpr value_type id() { return 0; }
  constexpr value_type operator()(const value_type &a, const value_type &b) {
    return a + b;
  }
  constexpr value_type inv(const value_type &x) { return -x; }
};

template <typename CommutativeGroup> class WeightedUnionFind {
  using T = typename CommutativeGroup::value_type;
  CommutativeGroup g;

  vector<int> tree, rank;
  vector<T> diffWeight;

public:
  WeightedUnionFind(const int &n) : tree(n), rank(n), diffWeight(n, g.id()) {
    iota(all(tree), 0);
  }

  int root(const int &x) {
    if (tree[x] == x) {
      return x;
    } else {
      int r = root(tree[x]);
      diffWeight[x] += diffWeight[tree[x]];
      return tree[x] = r;
    }
  }

  T weight(const int &x) {
    root(x);
    return diffWeight[x];
  }

  bool same(const int &a, const int &b) { return root(a) == root(b); }

  bool unite(int a, int b, T w) {
    w = g(w, g(weight(a), g.inv(weight(b))));
    a = root(a), b = root(b);
    if (a == b)
      return false;
    if (rank[a] < rank[b])
      swap(a, b), w = g.inv(w);
    tree[b] = a;
    diffWeight[b] = w;
    if (rank[a] == rank[b])
      ++rank[a];
    return true;
  }

  T diff(int x, int y) { return g(weight(y), g.inv(weight(x))); }
};

int main() {
  int n, m;
  while (cin >> n >> m, n || m) {
    WeightedUnionFind<sumGroup<int>> uf(n);
    rep(i, m) {
      char c;
      cin >> c;
      if (c == '!') {
        int a, b, w;
        cin >> a >> b >> w;
        uf.unite(--a, --b, w);
      } else {
        int a, b;
        cin >> a >> b;
        a--;
        b--;
        if (not uf.same(a, b)) {
          cout << "UNKNOWN" << endl;
        } else {
          cout << uf.diff(a, b) << endl;
        }
      }
    }
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define range(i, a, b) for (int i = (a); i < (b); i++)
#define rep(i, b) for (int i = 0; i < (b); i++)
#define all(a) (a).begin(), (a).end()
#define show(x) cerr << #x << " = " << (x) << endl;
const int INF = 1e8;
using namespace std;

static const int MAX_V = 100005;
static const int MAX_M = 100005;

class Topological {
public:
  vector<int> g[MAX_V]; //??°??????
  vector<int> tp;       //????????????????????????????????????

  void bfs(int s, int indeg[MAX_V], bool used[MAX_V]) {
    queue<int> q;

    q.push(s);
    used[s] = true;

    while (not q.empty()) {
      int u = q.front();
      q.pop();
      tp.emplace_back(u);
      rep(i, g[u].size()) {
        int v = g[u][i];
        indeg[v]--;
        if (indeg[v] == 0 && not used[v]) {
          used[v] = true;
          q.push(v);
        }
      }
    }
  }

  //??°??????????????????????????´??????0?????????
  bool topologicalSort(int v) {
    int indeg[MAX_V]; //??\\?¬???°
    bool used[MAX_V];
    memset(indeg, 0, sizeof(indeg));
    memset(used, 0, sizeof(used));

    rep(i, v) rep(j, g[i].size()) indeg[g[i][j]]++;
    rep(i, v) if (indeg[i] == 0 && not used[i]) bfs(i, indeg, used);

    for (auto it : tp)
      cout << it << endl;

    if (tp.size() == v)
      return true;
    else
      return false;
  }
};

class Union {
public:
  int par[MAX_M];   //???
  int depth[MAX_M]; //??¨?????±???

  void init(int n) {
    rep(i, n) {
      par[i] = i;
      depth[i] = 0;
    }
  }

  int find(int x) {
    if (par[x] == x) {
      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;

    if (depth[x] < depth[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (depth[x] == depth[y])
        depth[x]++;
    }
  }

  bool same(int x, int y) { return find(x) == find(y); }
};

class Edge {
public:
  int to, cost;
  Edge(int to, int cost) : to(to), cost(cost) {}
};

class Node {
public:
  int dis;
  bool isUsed;
  int label;
  Node() {
    this->dis = INF;
    this->isUsed = 0;
  }
};

typedef vector<vector<Edge>> AdjList;

void dijkstra(AdjList g, int start, int n, vector<Node> &node) {
  priority_queue<int, vector<pair<int, int>>, greater<pair<int, int>>> q;

  q.push(make_pair(0, start));
  node[start].dis = 0;

  pair<int, int> u;
  while (not q.empty()) {
    u = q.top();
    q.pop();
    int current = u.second;
    node[current].isUsed = 1;

    rep(i, g[current].size()) {
      int next = g[current][i].to;
      if (node[next].isUsed == 0) {
        if (node[next].dis > node[current].dis + g[current][i].cost) {
          node[next].dis = node[current].dis + g[current][i].cost;
          q.push(make_pair(node[next].dis, next));
        }
      }
    }
  }
}

int main() {
  int n, q;
  while (cin >> n >> q, n) {
    int a[MAX_M], b[MAX_M], c[MAX_M];
    char com[MAX_M];

    AdjList g(n);
    vector<int> in(n, 0);
    rep(i, q) {
      cin >> com[i];
      if (com[i] == '!') {
        cin >> a[i] >> b[i] >> c[i];
        a[i]--;
        b[i]--;
        g[b[i]].emplace_back(Edge{a[i], c[i]});
        in[a[i]]++;
      } else {
        cin >> a[i] >> b[i];
        a[i]--;
        b[i]--;
      }
    }

    vector<Node> node(n);
    rep(i, n) {
      if (in[i] == 0) { //??\\?¬???°???0
        dijkstra(g, i, n, node);
      }
    }

    // for(auto i : node){ cout << i.dis << endl; }

    Union uni;
    uni.init(n);
    rep(i, q) {
      if (com[i] == '!') {
        uni.unite(a[i], b[i]);
      } else {
        if (uni.same(a[i], b[i])) {
          // show(a[i])
          cout << node[a[i]].dis - node[b[i]].dis << endl;
        } else {
          cout << "UNKNOWN" << endl;
        }
      }
    }
  }
}