[codeforces] F. Make It Connected

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : F. Make It Connected
Source            : codeforces
Category          : Graph Theory, Data Structure
Algorithm         : Minimum Spanning Tree, Kruskal Algorithm, Segment Tree
Verdict           : Accepted

#include "bits/stdc++.h"
 
using namespace std;
 
#define ll                    long long int
#define inf                   (ll)1e17
#define min3(a, b, c)         min(a, min(b, c))
#define min4(a, b, c, d)      min(a, min(b, min(c, d)))
 
static const int maxn = 2e5 + 5;
 
int father[maxn];
 
inline void init()
{
      for (int i = 1; i < maxn; i++) father[i] = i;
}
 
inline int findRep(int r)
{
      if (r == father[r]) return r;
      return father[r] = findRep(father[r]);
}
 
struct segmentTree
{
      ll minOne, minTwo;
      int minOnePosition, minTwoPosition;
      inline void assignLeaf(ll val, int pos, int ff)
      {
            minOne = val;
            minOnePosition = pos;
            minTwo = inf;
            minTwoPosition = -1;
      }
      inline void marge(const segmentTree &lft, const segmentTree &rgt)
      {
            ll cst[5];
            int pos[5];
            cst[1] = lft.minOne; pos[1] = lft.minOnePosition;
            cst[2] = lft.minTwo; pos[2] = lft.minTwoPosition;
            cst[3] = rgt.minOne; pos[3] = rgt.minOnePosition;
            cst[4] = rgt.minTwo; pos[4] = rgt.minTwoPosition;
            minOne = minTwo = inf;
            minOnePosition = minTwoPosition = -1;
            int fone = -1;
            int ftwo = -1;
            int take = -1;
            for (int i = 1; i <= 4; i++)
            {
                  if (cst[i] < minOne)
                  {
                        take = i;
                        minOne = cst[i];
                        minOnePosition = pos[i];
                        fone = findRep(pos[i]);
                  }
            }
            for (int i = 1; i <= 4; i++)
            {
                  if (i == take) continue;
                  int p = findRep(pos[i]);
                  if (cst[i] <= minTwo && p != fone)
                  {
                        minTwo = cst[i];
                        minTwoPosition = pos[i];
                        ftwo = p;
                  }
            }
      }
} Tree[maxn << 2];
 
 
struct node
{
      int u, v;
      ll w;
      node(int u = 0, int v = 0, ll w = 0) : u(u), v(v), w(w) {}
      inline friend bool operator < (const node A, const node B)
      {
            return A.w < B.w;
      }
};
 
using pii = pair <int , int>;
 
map <pii, ll> mapper;
ll cost[maxn];
vector <node> graph;
 
inline void build(int node, int a, int b)
{
      if (a == b)
      {
            Tree[node].assignLeaf(cost[a], a, father[a]);
            return;
      }
      int lft = node << 1;
      int rgt = lft | 1;
      int mid = (a + b) >> 1;
      build(lft, a, mid);
      build(rgt, mid+1, b);
      Tree[node].marge(Tree[lft], Tree[rgt]);
}
 
inline void update(int node, int a, int b, int pos, int p)
{
      if (a > b or a > pos or b < pos) return;
      if (a >= pos && b <= pos)
      {
            Tree[node].assignLeaf(cost[pos], pos, p);
            return;
      }
      int lft = node << 1;
      int rgt = lft | 1;
      int mid = (a + b) >> 1;
      update(lft, a, mid, pos, p);
      update(rgt, mid+1, b, pos, p);
      Tree[node].marge(Tree[lft], Tree[rgt]);
}
 
inline segmentTree query(int node, int a, int b, int i, int j)
{
      if (a == i && b == j) return Tree[node];
      int lft = node << 1;
      int rgt = lft | 1;
      int mid = (a + b) >> 1;
      segmentTree p, q, res;
      if (j <= mid)
      {
            res = query(node, a, mid, i, j);
      }
      else if (i > mid)
      {
            res = query(node, mid+1, b, i, j);
      }
      else
      {
            p = query(node, a, mid, i, mid);
            q = query(node, mid+1, b, mid+1, j);
            res.marge(p, q);
      }
      return res;
}
 
int main()
{
      //freopen("in.txt", "r", stdin);
      int n, m;
      scanf("%d %d", &n, &m);
      for (int i = 1; i <= n; i++) scanf("%lld", &cost[i]);
      for (int i = 1; i <= m; i++)
      {
            int u, v;
            ll w;
            scanf("%d %d %lld", &u, &v, &w);
            if (cost[u] + cost[v] > w) graph.emplace_back(u, v, w);
      }
      init();
      build(1, 1, n);
      int taken = 0;
      while (taken < n-1)
      {
            segmentTree qry = query(1, 1, n, 1, n);
            int u = qry.minOnePosition;
            int v = qry.minTwoPosition;
            int p = findRep(qry.minOnePosition);
            int q = findRep(qry.minTwoPosition);
            ll costFromST = inf;
            if (p != q)
            {
                  costFromST = qry.minOne + qry.minTwo;
            }
            if (costFromST == inf)
            {
                  break;
            }
            else
            {
                  father[q] = p;
                  update(1, 1, n, u, p);
                  update(1, 1, n, v, p);
                  graph.emplace_back(u, v, costFromST);
                  taken++;
            }
 
      }
      init();
      sort(graph.begin(), graph.end());
      taken = 0;
      ll tcost = 0;
      for (node g : graph)
      {
            int u = g.u;
            int v = g.v;
            ll w = g.w;
            int p = findRep(u);
            int q = findRep(v);
            if (p != q)
            {
                  father[q] = p;
                  tcost += w;
                  taken++;
            }
            if (taken == n-1) break;
      }
      printf("%lld", tcost);
}