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

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

const long long inf = 0x3f3f3f3f3f3f3f3fLL;
typedef pair<int, int> PII;
typedef unsigned int uint;
typedef unsigned long long uLL;

const int N = 1e5 + 10;

int n, m;

vector<PII> adj[N], radj[N];

uLL cnt1[N], cnt2[N];
long long dis1[N], dis2[N];

struct node {
  int u, v;
  long long dis;
  node(int u = 0, int v = 0, long long dis = 0) : u(u), v(v), dis(dis) {}
  bool operator<(const node &rhs) const { return dis > rhs.dis; }
};

vector<node> E;

int main() {
  int n, m;
  scanf("%d%d", &n, &m);
  for (int i = 0; i < m; i++) {
    int u, v, w;
    scanf("%d%d%d", &u, &v, &w);
    adj[u].push_back(make_pair(v, w));
    radj[v].push_back(make_pair(u, w));
    E.push_back(node(u, v, w));
  }

  memset(dis1, 0x3f, sizeof dis1);
  memset(dis2, 0x3f, sizeof dis2);
  // count all shortest route from 1 to other nodes
  priority_queue<node> pq;
  pq.push(node(0, 1, 0));
  cnt1[0] = 1;
  while (!pq.empty()) {
    node cur = pq.top();
    pq.pop();
    if (dis1[cur.v] < cur.dis)
      continue;
    if (dis1[cur.v] == cur.dis) {
      cnt1[cur.v] += cnt1[cur.u];
      continue;
    }
    dis1[cur.v] = cur.dis;
    cnt1[cur.v] = cnt1[cur.u];
    for (int i = 0; i < adj[cur.v].size(); i++) {
      long long d = dis1[cur.v] + adj[cur.v][i].second;
      if (dis1[adj[cur.v][i].first] >= d)
        pq.push(node(cur.v, adj[cur.v][i].first, d));
    }
  }

  assert(pq.empty() == true);
  // count all shortest route from 2 to other nodes
  pq.push(node(0, 2, 0));
  cnt2[0] = 1;
  while (!pq.empty()) {
    node cur = pq.top();
    pq.pop();
    if (dis2[cur.v] < cur.dis)
      continue;
    if (dis2[cur.v] == cur.dis) {
      cnt2[cur.v] += cnt2[cur.u];
      continue;
    }
    dis2[cur.v] = cur.dis;
    cnt2[cur.v] = cnt2[cur.u];
    for (int i = 0; i < radj[cur.v].size(); i++) {
      long long d = dis2[cur.v] + radj[cur.v][i].second;
      if (dis2[radj[cur.v][i].first] >= d)
        pq.push(node(cur.v, radj[cur.v][i].first, d));
    }
  }

  long long shortdis = dis1[2];
  for (int i = 0; i < E.size(); i++) {
    if (i == 5)
      i = i;
    node cur = E[i];
    long long newdis = dis1[cur.v] + dis2[cur.u] + cur.dis;
    if (shortdis > dis1[cur.v] + dis2[cur.u] + cur.dis) {
      puts("HAPPY");
    } else if (shortdis == newdis) {
      puts("SOSO");
    } else if (shortdis < dis1[cur.u] + dis2[cur.v] + cur.dis) {
      puts("SOSO");
    } else if (cnt1[cur.u] * cnt2[cur.v] == cnt1[2]) {
      puts("SAD");
    } else {
      puts("SOSO");
    }
  }

  return 0;
}
/*
4 5
1 3 5
3 4 6
4 2 7
2 1 18
3 2 13
*/


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_policy.hpp>

using namespace std;
// using namespace __gnu_pbds;

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<int, int> PII;
typedef pair<long long, long long> PLL;

#define si(a) scanf("%d", &a)
#define sii(a, b) scanf("%d %d", &a, &b)
#define siii(a, b, c) scanf("%d %d %d", &a, &b, &c)

#define sl(a) scanf("%lld", &a)
#define sll(a, b) scanf("%lld %lld", &a, &b)
#define slll(a, b, c) scanf("%lld %lld %lld", &a, &b, &c)

#define un(x) x.erase(unique(all(x)), x.end())
#define xx first
#define yy second
#define pb push_back
#define mp make_pair
#define all(v) v.begin(), v.end()
#define D(x) cerr << #x " = " << x << '\\n'
#define DBG cerr << "Hi!" << '\\n'

#define CLR(a) memset(a, 0, sizeof(a))
#define SET(a) memset(a, -1, sizeof(a))

#define eps 1e-9
#define PI acos(-1.0)

int setBit(int n, int pos) { return n = n | (1 << pos); }
int resetBit(int n, int pos) { return n = n & ~(1 << pos); }
bool checkBit(ll n, ll pos) { return (bool)(n & (1LL << pos)); }

// int fx[] = {+0, +0, +1, -1, -1, +1, -1, +1};
// int fy[] = {-1, +1, +0, +0, +1, +1, -1, -1}; //Four & Eight Direction

/******************************************************************************************/

const int MAX = 100010;
const ll INF = 2000000000000LL;

vector<int> g[MAX], rg[MAX], ed[MAX], bucket[MAX];

int sdom[MAX], par[MAX], dom[MAX], dsu[MAX], label[MAX];
int arr[MAX], rev[MAX], Time, n, source;

void init(int _n, int _source) {
  Time = 0;
  n = _n;
  source = _source;
  for (int i = 1; i <= n; i++) {
    g[i].clear(), rg[i].clear(), ed[i].clear(), bucket[i].clear();
    arr[i] = sdom[i] = par[i] = dom[i] = dsu[i] = label[i] = rev[i] = 0;
  }
}

void dfs(int u) {
  Time++;
  arr[u] = Time;
  rev[Time] = u;
  label[Time] = Time;
  sdom[Time] = Time;
  dsu[Time] = Time;
  int i, w;
  for (i = 0; i < g[u].size(); i++) {
    w = g[u][i];
    if (!arr[w]) {
      dfs(w);
      par[arr[w]] = arr[u];
    }
    rg[arr[w]].pb(arr[u]);
  }
}

inline int Find(int u, int x = 0) {
  if (u == dsu[u])
    return x ? -1 : u;
  int v = Find(dsu[u], x + 1);
  if (v < 0)
    return u;
  if (sdom[label[dsu[u]]] < sdom[label[u]])
    label[u] = label[dsu[u]];
  dsu[u] = v;
  return x ? v : label[u];
}

inline void Union(int u, int v) { dsu[v] = u; }

void build() {
  dfs(source);
  for (int i = n; i >= 1; i--) {
    for (int j = 0; j < rg[i].size(); j++)
      sdom[i] = min(sdom[i], sdom[Find(rg[i][j])]);
    if (i > 1)
      bucket[sdom[i]].pb(i);
    for (int j = 0; j < bucket[i].size(); j++) {
      int w = bucket[i][j], v = Find(w);
      if (sdom[v] == sdom[w])
        dom[w] = sdom[w];
      else
        dom[w] = v;
    }
    if (i > 1)
      Union(par[i], i);

    for (int i = 2; i <= n; i++) {
      if (dom[i] != sdom[i])
        dom[i] = dom[dom[i]];

      if (rev[i] == 0 or rev[dom[i]] == 0)
        continue;

      ed[rev[i]].pb(rev[dom[i]]);
      ed[rev[dom[i]]].pb(rev[i]);
    }
  }
}

vector<int> G[2][MAX];
vector<int> co[2][MAX];
ll dis[2][MAX];
bool vis[MAX];

struct node {
  int city;
  ll cost;
};

bool operator<(node a, node b) { return a.cost > b.cost; }

void dijkstra(int s, int n, int id, int d) {
  int i, x, u, v, c;
  node a, b;
  for (int i = 1; i <= n; i++)
    dis[id][i] = INF;
  dis[id][s] = 0;
  a = {s, 0};
  priority_queue<node> Q;
  Q.push(a);

  while (!Q.empty()) {
    a = Q.top();
    Q.pop();
    u = a.city;
    if (!vis[u]) {
      vis[u] = true;
      for (int i = 0; i < G[id][u].size(); i++) {
        v = G[id][u][i];
        c = co[id][u][i];
        if (dis[id][v] > dis[id][u] + c) {
          dis[id][v] = dis[id][u] + c;
          b = {v, dis[id][v]};
          if (v != d)
            Q.push(b);
        }
      }
    }
  }
}

int lg;
int L[MAX], P[MAX][22];

void dfss(int s, int par, int l) {
  int i, x;
  L[s] = l;
  for (int i = 0; i < ed[s].size(); i++) {
    x = ed[s][i];
    if (x != par) {
      P[x][0] = s;
      dfss(x, s, l + 1);
    }
  }
}

void lca_build(int n, int root) {
  SET(P);
  dfss(root, -1, 0);
  lg = (log(n) / log(2.0)) + 2;

  int i, j;
  for (int j = 1; (1 << j) <= n; j++) {
    for (i = 1; i <= n; i++) {
      if (P[i][j - 1] != -1)
        P[i][j] = P[P[i][j - 1]][j - 1];
    }
  }
}

inline int lca_query(int x, int y) {
  if (L[x] < L[y])
    swap(x, y);
  int i, j;
  for (i = lg; i >= 0; i--) {
    if (L[x] - (1 << i) >= L[y])
      x = P[x][i];
  }

  if (x == y)
    return x;
  for (i = lg; i >= 0; i--) {
    if (P[x][i] != -1 && P[x][i] != P[y][i]) {
      x = P[x][i];
      y = P[y][i];
    }
  }
  return P[x][0];
}

struct Edge {
  int u, v, w;
} edges[MAX];

map<PII, int> Mn;
map<PII, int> Cnt;

bool ase[MAX];

int main() {
  //    freopen("in.txt", "r", stdin);
  //    freopen("out.txt", "w", stdout);

  int m, u, v, w;
  sii(n, m);

  init(n, 1);

  for (int i = 1; i <= m; i++) {
    siii(u, v, w);
    edges[i] = {u, v, w};

    if (Mn.find({u, v}) == Mn.end()) {
      Mn[{u, v}] = w;
      Cnt[{u, v}] = 1;
    } else {
      if (w == Mn[{u, v}])
        Cnt[{u, v}]++;
      else if (w < Mn[{u, v}]) {
        Mn[{u, v}] = w;
        Cnt[{u, v}] = 1;
      }
    }
    G[0][u].push_back(v);
    co[0][u].push_back(w);

    G[1][v].push_back(u);
    co[1][v].push_back(w);
  }

  dijkstra(1, n, 0, 2);
  dijkstra(2, n, 1, 1);

  for (int i = 1; i <= m; i++) {
    u = edges[i].u;
    v = edges[i].v;
    w = edges[i].w;
    if (dis[0][u] + w + dis[1][v] == dis[0][2]) {
      g[u].pb(v);
      ase[i] = true;
    }
  }

  build();
  lca_build(n, 1);
  for (int i = 1; i <= m; i++) {
    u = edges[i].u;
    v = edges[i].v;
    w = edges[i].w;

    if (dis[0][v] + w + dis[1][u] < dis[0][2]) {
      puts("HAPPY");
      continue;
    }

    if (!ase[i]) {
      puts("SOSO");
      continue;
    }

    if (Cnt[{u, v}] >= 2) {
      puts("SOSO");
      continue;
    }

    if (lca_query(u, 2) == u && lca_query(v, 2) == v) {
      puts("SAD");
      continue;
    } else {
      puts("SOSO");
      continue;
    }
  }

  return 0;
}