HackerRank The Value of Friendship problem solution

In this HackerRank The Value of Friendship problem solution You are given q queries, where each query is in the form of an unordered list of m distinct direct friendships between n students. For each query, find the maximum value of total among all possible orderings of formed friendships and print it on a new line.

Problem solution in Python.

def grow(n):
  res = 0
  for i in range(1, n):
    res += i*(i+1)
    #print("-->", i, res)
  return res
  
def mark(x):
  z = x
  while z: z = z[0]
  if x is not z: x[0] = z
  return z

def mark2(x):
  z = x
  while z and type(z) == list: z = z[0]
  if x is not z: x[0] = z
  return z

  
T = int(input())
for t in range(T):
  n, m = map(int, input().split())
  work = [[] for i in range(n)] 
  lost = 0
  for i in range(m):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    mu = mark(work[u])
    mv = mark(work[v])
    if mu is not mv:
      mu.append(mv)
    else:
      lost += 1
      
  no = 1
  dic = {}
  for i in range(n):
    z = mark2(work[i])
    if not z:
      z.append(no)
      z = no
      no += 1
      dic[z] = 1
    else:
      dic[z] += 1

  res = 0
  cur = 0    
  for v in sorted(dic.values(), reverse = True):
    res += grow(v) + (v-1) * cur
    cur += v*(v-1)
    #print(v, cur, res) 
    
  res += lost * cur
  print(res)  

Problem solution in Java.

import java.util.*;

public class Solution {

    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int t = in.nextInt();
        for(int a0 = 0; a0 < t; a0++){
            int n = in.nextInt();
            int m = in.nextInt();

            DisjointSets friendGroups = new DisjointSets(n);
            
            for(int a1 = 0; a1 < m; a1++){
                int x = in.nextInt();
                int y = in.nextInt();
                // your code goes here
                friendGroups.union(x - 1, y - 1);
            }
            
            List<Integer> groupSizes = friendGroups.setSizes();
            Collections.sort(groupSizes, Comparator.reverseOrder());

            long[] extra = new long[m];
            int idx = 0;
            for (int size : groupSizes) {
                for (int i = 1; i < size; i++) {
                    extra[idx++] = i;
                }
            }
            
            long totalFriendship = 0;
            long currentFriendship = 0;
            for (int i = 0; i < m; i++) {
                currentFriendship += 2L * extra[i];
                totalFriendship += currentFriendship;
            }
            System.out.println(totalFriendship);
        }
    }
    
    private static class DisjointSets {
        private int[] parent;
        private int[] size;
        
        public DisjointSets(int n) {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++) {
                parent[i] = i;
                size[i] = 1;
            }
        }
        
        public void union(int a, int b) {
            int p = find(a);
            int q = find(b);
            if (p == q)
                return;
            if (size[p] < size[q]) {
                parent[p] = q;
                size[q] += size[p];
            }
            else {
                parent[q] = p;
                size[p] += size[q];
            }
        }
        
        public int find(int x) {
            while (x != parent[x]) {
                parent[x] = parent[parent[x]];
                x = parent[x];
            }
            return x;
        }
        
        public List<Integer> setSizes() {
            Set<Integer> sets = new HashSet<>();
            for (int x : parent)
                sets.add(find(x));
            List<Integer> sizes = new ArrayList<>();
            for (int x : sets)
                sizes.add(size[x]);
            return sizes;
        }
    }
}

Problem solution in C++.

#include <map>
#include <set>
#include <list>
#include <cmath>
#include <ctime>
#include <deque>
#include <queue>
#include <stack>
#include <string>
#include <bitset>
#include <cstdio>
#include <limits>
#include <vector>
#include <climits>
#include <cstring>
#include <cstdlib>
#include <fstream>
#include <numeric>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <unordered_map>

using namespace std;

vector<int> parent, setSz;


int find(int x)
{
	if (parent[x] == x)
		return x;
	return parent[x] = find(parent[x]);
}

int setunion(int a, int b)
{
	a = find(a);
	b = find(b);
	if (a == b)
		return setSz[a];
	if (setSz[a] < setSz[b])
		swap(a, b);
	parent[b] = a;
	setSz[a] += setSz[b];
	return setSz[a];
}

#define lint long long int


int main() {


	int t;
	cin >> t;
	for (int a0 = 0; a0 < t; a0++) {
		int n;
		int m;
		cin >> n >> m;

		parent.clear();
		parent.resize(n + 1);
		for (int i = 1; i <= n; i++)
			parent[i] = i;
		setSz.clear();
		setSz.resize(n + 1, 1);

		lint curFr = 0;
		lint res = 0;

		int same = 0;
	
		for (int a1 = 0; a1 < m; a1++) {
			int x;
			int y;
			cin >> x >> y;
			if (find(x) == find(y))
				same++;
			setunion(x, y);
			//lint oldSz1 = setSz[find(x)];
			//lint oldSz2 = setSz[find(y)];
			//lint newSz = setunion(x, y);
			//curFr += (newSz * (newSz - 1LL)) - (oldSz1 * (oldSz1 - 1LL)) - (oldSz2 * (oldSz2 - 1LL));
			//res += curFr;

			//cout << newSz << " " << curFr << endl;
		}
		set<int> clusters;
		vector<int> clusterSz;
		for (int i = 1; i <= n; i++)
		{
			int x = find(i);
			if (clusters.find(x) != clusters.end())
				continue;
			clusterSz.push_back(setSz[x]);
			clusters.insert(x);
		}
		sort(clusterSz.begin(), clusterSz.end());


		for (int i = clusterSz.size() - 1; i >= 0; i--)
		{
			for (int j = 2; j <= clusterSz[i]; j++)
			{
				curFr += j * (j - 1) - (j - 1) * (j - 2);
				res += curFr;
			}
		}
		res += curFr * same;


		cout << res << endl;
	}

	return 0;
}

Problem solution in C.

#include <stdio.h>
#include <stdint.h>
#include <inttypes.h>
#include <string.h>
#include <stdlib.h>

int n, m;
int graph[500002][2];  // [0]-neighbor node, [1]-next
int gend;
int nodeedge[100001];
int used[100001];

int bfs(int v, int *edgen) {
  int q[100001];
  int i, j, qe=1;
  q[0]=v;
  used[v]=1;
  *edgen=nodeedge[v];
  for (i=0; i<qe; i++) {
    for (j=q[i]; j && (v=graph[j][0]); j=graph[j][1]) {
      if (!used[v]) {
        q[qe++]=v;
        used[v]=1;
        *edgen+=nodeedge[v];
      }
    }
  }
  return qe;
}

static int cmp(const void *p1, const void *p2) {
  return *((const int *)p2)-*((const int *)p1);
}

uint64_t solve(void) {
  uint64_t total=0, tt=0, t;
  int noden[100001], edgen, rededgen=0;
  int i, nni=0;
  for (i=1; i<=n; i++) {
    if (used[i])
      continue;
    noden[nni]=bfs(i, &edgen);
    if (edgen)
      rededgen+=edgen-noden[nni++]+1;
  }
  qsort(noden, nni, sizeof(int), cmp);
  for (i=0; i<nni; i++) {
    t=noden[i];
    total+=(t*(t+1)*(t-1))/3+(t-1)*tt;
    tt+=t*(t-1);
  }
  total+=tt*rededgen;
  return total;
}

int main(void) {
  int i, q, u, v;
  scanf("%dn", &q);
  while (q--) {
    scanf("%d%dn", &n, &m);
    gend=n+1;
    for (i=0; i<m; i++) {
      scanf("%d%dn", &u, &v);
      if (graph[u][0]) {
        graph[gend][0]=graph[u][0];
        graph[gend][1]=graph[u][1];
        graph[u][0]=v;
        graph[u][1]=gend++;
      } else {
        graph[u][0]=v;
      }
      if (graph[v][0]) {
        graph[gend][0]=graph[v][0];
        graph[gend][1]=graph[v][1];
        graph[v][0]=u;
        graph[v][1]=gend++;
      } else {
        graph[v][0]=u;
      }
      nodeedge[u]++;
    }
    printf("%"PRIu64"n", solve());
    memset(graph, 0, sizeof(graph));
    memset(nodeedge, 0, sizeof(nodeedge));
    memset(used, 0, sizeof(used));
    gend=0;
  }
  return 0;
}