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

#include <bits/stdc++.h>
#define int long long
#define N 100010
using namespace std;
const int INF = 1LL << 55;
const int mod = (1e9) + 7;
const double EPS = 1e-8;
const double PI = 6.0 * asin(0.5);
template <class T> T Max(T &a, T b) { return a = max(a, b); }
template <class T> T Min(T &a, T b) { return a = min(a, b); }

#define rank Asdfiasofsad
class UF {
public:
  int V;
  vector<int> par, rank;
  UF() {}
  UF(int V) : V(V), par(V), rank(V, 0) {
    for (int i = 0; i < V; i++)
      par[i] = i;
  }

  int find(int x) {
    assert(x < V);
    if (par[x] == x)
      return x;
    return par[x] = find(par[x]);
  }

  void unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return;
    if (rank[x] < rank[y])
      par[x] = y;
    else {
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }

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

class RMQ2 {
public:
  typedef long long ll;
  struct data {
    bool type; // 0 - empty   , 1 - update
    ll value;
  };
  ll n;
  vector<ll> dat;
  vector<data> td;

  RMQ2() { n = -1; }
  RMQ2(int n_) {
    n = 1;
    while (n < n_)
      n *= 2;
    td.resize(2 * n - 1, (data){0, 0});
    dat.resize(2 * n - 1, INF);
  }

  //[a,b)?????????x????????´???update(a,b,x)
  ll update(int a, int b, int x, bool flg = true, int k = 0, int l = 0,
            int r = -1) {
    if (r == -1)
      r = n, assert(a < n && b <= n);
    if (r <= a || b <= l)
      return flg ? dat[k] : INF;
    if (a <= l && r <= b) {
      if (flg == true) {
        td[k] = (data){1, x};
        dat[k] = x;
      }
      return dat[k];
    }

    if (td[k].type) {
      dat[k * 2 + 1] = dat[k * 2 + 2] = td[k].value;
      td[k * 2 + 1] = td[k * 2 + 2] = (data){1, td[k].value};
      td[k].type = 0;
    }

    ll vl = update(a, b, x, flg, k * 2 + 1, l, (l + r) / 2);
    ll vr = update(a, b, x, flg, k * 2 + 2, (l + r) / 2, r);
    if (flg == true)
      dat[k] = min(vl, vr);
    return min(vl, vr);
  }

  //[a,b)???????°????????????????find(a,b);
  ll find(int a, int b) { return update(a, b, 0, false); };
};

/*???HL????§£???O(V+E)*/
class HLD {
public:
  int V;                 //???????????°
  int root;              //??????????????????
  vector<vector<int>> G; // Grapth
  vector<int> par;       // parent: par[??????????????????]->??????????????????
  vector<int> Hchild; //???????????????????????????Heavy Child???????????????:
                      //Hedge[??????????????????]->??????????????????
  vector<int> number; //???????????????????????????????????????????????????:
                      //number[??????????????????]->????????´??????????????????
  vector<int> id;     //??????????????´????????°???????????????????????????:
                      //id[????????´??????????????????]->??????????????????
  vector<int> group;  //???????????????????±?????????°???????????????:
                      //group[??????????????????]->??????????????????
  vector<int>
      depth; //?????????????????????????????±???: depth[??????????????????]
  bool ok;   // build()???????????§?????????

  HLD() { root = V = -1; }
  HLD(int V)
      : V(V), root(0), G(V), par(V, -1), Hchild(V, -1), number(V, -1),
        id(V, -1), group(V, -1), depth(V, 0), ok(0) {}

  void add_edge(int a, int b) {
    assert(a < V && b < V);
    G[a].push_back(b);
    G[b].push_back(a);
  }

  void build(int root = 0) {
    this->root = root;
    ok = 1;
    weightParent(); // weight and parent
    int num = 0;
    numbering(num);
    depthGroup();

    if (0) {
      cout << "par:";
      for (int a : par)
        cout << a << " ";
      cout << endl;
      cout << "group";
      for (int a : group)
        cout << a << " ";
      cout << endl;

      for (int i = 0; i < V; i++)
        for (int to : G[i])
          if (i < to)
            cout << i << "[" << number[i] << "] " << to << "[" << number[to]
                 << "]" << endl;
    }
  }

  int weightParent(int pos = -1, int pre = -1) {
    if (pos == -1)
      pos = root;
    par[pos] = pre;
    int mx = -1, weight = 1;
    for (int to : G[pos]) {
      if (to == pre)
        continue;
      int cweight = weightParent(to, pos);
      weight += cweight;
      if (mx < cweight)
        Hchild[pos] = to, mx = cweight;
    }
    return weight;
  }

  void depthGroup(int pos = -1, int pre = -1, int dep = 0, int g = -1) {
    if (pos == -1)
      pos = g = root;
    depth[pos] = dep;
    group[pos] = g;
    for (int to : G[pos])
      if (to != pre)
        depthGroup(to, pos, dep + 1, Hchild[pos] == to ? g : to);
  }

  void numbering(int &cnt, int pos = -1, int pre = -1) {
    if (pos == -1)
      pos = root;
    number[pos] = cnt++;
    id[number[pos]] = pos;
    if (Hchild[pos] >= 0)
      numbering(cnt, Hchild[pos], pos);
    for (int to : G[pos])
      if (to != pre && to != Hchild[pos])
        numbering(cnt, to, pos);
  }

  typedef pair<int, int> P;
  vector<P>
      paths; //[u,v]????????????
             //????????????????????£???????????????????????????????§£?????????(?????§logN??????????????????????????????)
  vector<P> lightEdges; // paths????????¶???????????¨??????
  void buildPaths(int u, int v, int init = true) {
    if (init)
      paths.clear(), lightEdges.clear(), assert(ok);
    if (depth[group[u]] < depth[group[v]])
      swap(u, v); // ??±???????????????????????????????????????
    int nu = number[u], nv = number[v];
    if (same(u, v)) {
      paths.push_back(P(min(nu, nv), max(nu, nv)));
      return;
    }

    int gu = group[u];
    buildPaths(par[gu], v, false);

    paths.push_back(P(number[gu], number[u]));
    lightEdges.push_back(P(number[par[gu]], number[gu]));
  }

  pair<vector<P>, vector<P>> getPath(int u, int v) {
    buildPaths(u, v);
    return make_pair(paths, lightEdges);
  }

  /*(??????????????????,??????????????????)*/
  bool same(int a, int b) { return group[a] == group[b]; }
};

typedef pair<int, int> P;
typedef pair<P, P> PP;
int n, m;
vector<vector<int>> G;
vector<PP> edge;
vector<int> used;
HLD hld;
RMQ2 rmq;

int kuraskal() {
  sort(edge.begin(), edge.end());
  hld = HLD(n);
  used.resize(m, 0);
  UF U(n);
  int sum = 0, cnt = 0;
  for (int i = 0; i < m; i++) {
    int a = edge[i].second.first;
    int b = edge[i].second.second;
    int c = edge[i].first.first;
    if (U.same(a, b))
      continue;
    cnt++;
    U.unite(a, b);
    hld.add_edge(a, b);
    used[i] = 1;
    sum += c;
  }
  if (cnt != n - 1)
    return INF;
  return sum;
}

void build() {
  rmq = RMQ2(n);
  int root = 0;
  for (int i = 0; i < n; i++)
    if (hld.G[i].size() > hld.G[root].size())
      root = i;

  hld.build(rand() % n);

  for (int i = m - 1; i >= 0; i--) {
    if (used[i])
      continue;
    int a = edge[i].second.first;
    int b = edge[i].second.second;
    int c = edge[i].first.first;
    pair<vector<P>, vector<P>> path = hld.getPath(a, b);
    vector<P> &hpath = path.first, &ledge = path.second;
    if (hld.depth[a] > hld.depth[b])
      swap(a, b);
    hpath[0].first++;
    for (P p : hpath)
      rmq.update(p.first, p.second + 1, c);
    for (P e : ledge)
      rmq.update(e.second, e.second + 1, c);
  }
}

signed main() {
  cin >> n >> m;
  G.resize(n);
  edge.resize(m);
  for (int i = 0; i < m; i++) {
    int a, b, c;
    cin >> a >> b >> c;
    a--, b--;
    edge[i] = PP(P(c, i), P(a, b));
  }

  int mincost = kuraskal();
  build();
  vector<int> ans(m, mincost);
  for (int i = 0; i < m; i++) {
    if (!used[i])
      continue;
    int idx = edge[i].first.second;
    int a = edge[i].second.first;
    int b = edge[i].second.second;
    int c = edge[i].first.first;
    if (hld.depth[a] > hld.depth[b])
      swap(a, b);
    int c2 = rmq.find(hld.number[b], hld.number[b] + 1);
    ans[idx] = mincost - c + c2;
  }

  for (int a : ans)
    cout << (a >= INF ? -1 : a) << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define N 100011
#define MAX_V 100010
#define MAX_LOG_V 20
#define INF (1e9)
using namespace std;
vector<int> G[MAX_V];
int root;
int parent[MAX_LOG_V][MAX_V];
int depth[MAX_V];

void dfs(int v, int p, int d) {
  parent[0][v] = p;
  depth[v] = d;
  for (int i = 0; i < G[v].size(); i++)
    if (G[v][i] != p)
      dfs(G[v][i], v, d + 1);
}

void init(int V) {
  dfs(root, -1, 0);
  for (int k = 0; k + 1 < MAX_LOG_V; k++)
    for (int v = 0; v < V; v++)
      if (parent[k][v] < 0)
        parent[k + 1][v] = -1;
      else
        parent[k + 1][v] = parent[k][parent[k][v]];
}

int lca(int u, int v) {
  if (depth[u] > depth[v])
    swap(u, v);
  for (int k = 0; k < MAX_LOG_V; k++)
    if ((depth[v] - depth[u]) >> k & 1)
      v = parent[k][v];
  if (u == v)
    return u;

  for (int k = MAX_LOG_V - 1; k >= 0; k--)
    if (parent[k][u] != parent[k][v]) {
      u = parent[k][u];
      v = parent[k][v];
    }
  return parent[0][u];
}

typedef long long ll;
struct edge {
  int to, cost, idx;
};
typedef pair<int, int> P;
typedef pair<ll, P> PP;
vector<edge> g[N];
int MST[N] = {};
map<P, int> M;
int n, m;

void add_edge(int a, int b, int cost, int idx) {
  g[a].push_back((edge){b, cost, idx});
  g[b].push_back((edge){a, cost, idx});
}

ll prim() {
  memset(MST, -1, sizeof(MST));
  ll res = 0;
  priority_queue<PP, vector<PP>, greater<PP>> Q;
  Q.push(PP(0, P(0, 0)));
  while (!Q.empty()) {
    PP t = Q.top();
    Q.pop();
    ll cost = t.first;
    int pos = t.second.first, pre = t.second.second;
    if (MST[pos] != -1)
      continue;
    MST[pos] = pre;
    res += cost;
    for (int i = 0; i < g[pos].size(); i++) {
      int to = g[pos][i].to;
      int cost = g[pos][i].cost;
      if (MST[to] != -1)
        continue;
      Q.push(PP(cost, P(to, pos)));
    }
  }
  return res;
}

typedef priority_queue<P, vector<P>, greater<P>> PQ;
int ans[200001], dep = 0, pos[N], pre[N];

PQ dfs() {
  PQ res;
  res.push(P(INF, -1));
  for (int i = 0; i < g[pos[dep]].size(); i++) {
    int LCA = lca(pos[dep], g[pos[dep]][i].to);
    if (LCA == pos[dep] || g[pos[dep]][i].to == pre[dep])
      continue;
    res.push(P(g[pos[dep]][i].cost, LCA));
  }

  for (int i = 0; i < G[pos[dep]].size(); i++) {
    dep++;
    pos[dep] = G[pos[dep - 1]][i];
    pre[dep] = pos[dep - 1];
    PQ t = dfs();
    dep--;
    while (!t.empty()) {
      P r = t.top();
      t.pop();
      if (r.second != pos[dep])
        res.push(r);
    }
  }
  assert(res.size() < m + 100);
  if (pos[dep])
    ans[M[P(pos[dep], pre[dep])]] = res.top().first;
  return res;
}

int main() {
  cin >> n >> m;
  vector<PP> e(m);
  for (int i = 0, a, b, c; i < m; i++) {
    scanf("%d%d%d", &a, &b, &c);
    a--, b--;
    e[i] = PP(c, P(a, b));
    M[P(a, b)] = M[P(b, a)] = i;
    add_edge(a, b, c, i);
  }

  ll mncost = prim();
  int used[200011] = {}, flg = 0;
  for (int i = 1; i < n; i++) {
    flg |= MST[i] == -1;
    if (MST[i] == -1)
      continue;
    G[MST[i]].push_back(i);
    used[M[P(i, MST[i])]] = 1;
  }
  init(n);
  pos[0] = 0, pre[0] = -1;
  dfs();
  for (int i = 0; i < m; i++) {
    if (!used[i] && !flg)
      printf("%lld\\n", mncost);
    else if (ans[i] >= INF)
      printf("-1\\n");
    else
      printf("%lld\\n", mncost - e[i].first + ans[i]);
  }
  return 0;
}