A tech fest is buzzing with activity as employees from various companies aim to connect and collaborate. Each connection between employees X and Y is represented as “C X Y” which leads to a merger of their respective companies if they belong to different ones. Initially, N employees represent N companies. You need to implement a system that handles two types of queries efficiently:
- Connection Query (C X Y): Establishes a connection between employees X and Y. If they belong to different companies, their companies merge.
- Query (Q X): Retrieves the size of the company to which employee X belongs.
Examples:
Input: N = 2, Q = 3
- Q 1
- C 1 2
- Q 2
Output:
- 1
- 2
Explanation:
- Query1- Size of the company where Employee 1 belongs to
- Size is 1
- Query2- Connect companies where Employee 1 & Employee 2 belongs to
- Query3- Size of the company where Employee 2 belongs to
- Size is 2
Input: N = 3, Q = 3
- Q 1
- C 2 3
- Q 3
Output:
- 1
- 2
Approach: To solve the problem follow the below idea:
The idea is to use the Disjoint-Set Union(DSU) data structure.
The approach using Disjoint-Set Union:
- Let’s treat every employee as a connected component and maintain a Disjoint-Set Union(DSU) data structure.
- Initially, there are components where the i-th component consists of the i-th employee.
- The query for connecting two employees can be handled using the union operation of our Disjoint-Set Union data structure.
- Now to get the size of a component(company) we need to modify the union function a bit.
- For every component let’s store the size of that component. Initially, all the components are of size 1(as only 1 person is there). The size of a component gets updated when it merges with another component and the new size is equal to the sum of the sizes of the components that merged.
Implementation of the above approach:
C++14
// C++ code for the above approach:#include <iostream>#include <vector>using namespace std;
class UnionFind {
private:
// To store the parent of each
vector<int> parent;
// To store the size of the every company
vector<int> cnt;
public:
// Initialize the UnionFind class
UnionFind(int size)
: parent(size + 1)
, cnt(size + 1)
{
for (int i = 1; i <= size; ++i) {
parent[i] = i;
cnt[i] = 1;
}
}
// Find company of the employee
int find(int node)
{
// If their is only 1 employee in
// the company, return it
if (node == parent[node])
return node;
// Else return it's parent
return parent[node] = find(parent[node]);
}
// Return the size of the company
int totalNodesInSet(int node)
{
return cnt[find(node)];
}
// To connect companies where employee a
// and employee b belongs to
void Union(int a, int b)
{
// Find employee a's company/parent
a = find(a);
// Find employee b's company/parent
b = find(b);
// If both belongs to same company,
// no need to do anything
if (a == b)
return;
// Set parent/company of employee b to a,
// as both belongs to same company now
parent[b] = a;
// Recalculate the size of the company
// after merging
cnt[a] += cnt[b];
}
};// Drivers codeint main()
{ // n = no of employees & q = no of queries
int n, q;
n = 3, q = 3;
// Call Unionfind Class to initialize
// and build with n employees
UnionFind uf(n);
// Calculates the size of node 1
cout << uf.totalNodesInSet(1) << "\n";
// Joins components in which 2 and 3
// are present
uf.Union(2, 3);
// Calculates the size of node 3
cout << uf.totalNodesInSet(3) << "\n";
return 0;
} |
Java
// Java code for the above approach:import java.util.Arrays;
class UnionFind {
private int[] parent;
private int[] cnt;
public UnionFind(int size) {
// To store the parent of each
// company/employee
parent = new int[size + 1];
// To store the size of each company
cnt = new int[size + 1];
for (int i = 1; i <= size; ++i) {
parent[i] = i;
cnt[i] = 1;
}
}
public int find(int node) {
// If there is only 1 employee in
// the company, return it
if (node == parent[node]) {
return node;
}
// Else return its parent
return parent[node] = find(parent[node]);
}
public int totalNodesInSet(int node) {
// Return the size of the company
return cnt[find(node)];
}
public void union(int a, int b) {
// Find employee a's company/parent
a = find(a);
// Find employee b's company/parent
b = find(b);
// If both belong to the same company,
// no need to do anything
if (a == b) {
return;
}
// Set parent/company of employee b to a,
// as both belong to the same company now
parent[b] = a;
// Recalculate the size of the company
// after merging
cnt[a] += cnt[b];
}
}public class Main {
public static void main(String[] args) {
// n = no of employees & q = no of queries
int n = 3;
int q = 3;
// Call UnionFind Class to initialize
// and build with n employees
UnionFind uf = new UnionFind(n);
// Calculates the size of node 1
System.out.println(uf.totalNodesInSet(1));
// Joins components in which 2 and 3
// are present
uf.union(2, 3);
// Calculates the size of node 3
System.out.println(uf.totalNodesInSet(3));
}
}//this code is contributed by uttamdp_10 |
Python
n, queries = map(int, input().split())
# data structure to keep parent and size of each sub grouppeople = [{'parent': i, 'n': i, 'size': 1} for i in range(n)]
# return the index of the parentdef parent_index(n):
while people[n]['parent'] != people[n]['n']:
# path compression (works logically without this, but will
# fail last 3 tests without it)
people[n]['parent'] = people[people[n]['parent']]['parent']
n = people[n]['parent']
return n
for _ in range(queries):
q = input().split()
if q[0] == 'Q':
n = int(q[1])-1
# print the size of parent
print(people[parent_index(n)]['size'])
elif q[0] == 'M':
p1, p2 = parent_index(int(q[1])-1), parent_index(int(q[2])-1)
if p1 != p2:
# set the parent of p2 to be the parent of p1
people[p2]['parent'] = people[p1]['parent']
people[p1]['size'] += people[p2]['size']
|
C#
// C# code for the above approach:using System;
class UnionFind
{ private int[] parent;
private int[] cnt;
public UnionFind(int size)
{
// To store the parent of each
// company/employee
parent = new int[size + 1];
// To store the size of each company
cnt = new int[size + 1];
for (int i = 1; i <= size; ++i)
{
parent[i] = i;
cnt[i] = 1;
}
}
public int Find(int node)
{
// If there is only 1 employee in
// the company, return it
if (node == parent[node])
{
return node;
}
// Else return its parent
return parent[node] = Find(parent[node]);
}
public int TotalNodesInSet(int node)
{
// Return the size of the company
return cnt[Find(node)];
}
public void Union(int a, int b)
{
// Find employee a's company/parent
a = Find(a);
// Find employee b's company/parent
b = Find(b);
// If both belong to the same company,
// no need to do anything
if (a == b)
{
return;
}
// Set parent/company of employee b to a,
// as both belong to the same company now
parent[b] = a;
// Recalculate the size of the company
// after merging
cnt[a] += cnt[b];
}
}public class GFG
{ public static void Main(string[] args)
{
// n = no of employees & q = no of queries
int n = 3;
// Call UnionFind Class to initialize
// and build with n employees
UnionFind uf = new UnionFind(n);
// Calculates the size of node 1
Console.WriteLine(uf.TotalNodesInSet(1));
// Joins components in which 2 and 3
// are present
uf.Union(2, 3);
// Calculates the size of node 3
Console.WriteLine(uf.TotalNodesInSet(3));
}
}// This code is contributed by Sakshi |
Javascript
class GFG { constructor(size) {
// To store the parent of each node
this.parent = Array(size + 1);
// To store the size of every company
this.cnt = Array(size + 1);
for (let i = 1; i <= size; ++i) {
this.parent[i] = i;
this.cnt[i] = 1;
}
}
// Find the company of the employee
find(node) {
// If there is only 1 employee in
// the company return it
if (node === this.parent[node]) return node;
// Else return its parent and
// perform path compression
return (this.parent[node] = this.find(this.parent[node]));
}
// Return the size of the company
totalNodesInSet(node) {
return this.cnt[this.find(node)];
}
// Connect companies where employee a and
// employee b belong to
union(a, b) {
// Find the company of employee a
a = this.find(a);
// Find the company of employee b
b = this.find(b);
// If both belong to the same company
// no need to do anything
if (a === b) return;
// Set the parent/company of employee b to a as both
// belong to the same company now
this.parent[b] = a;
// Recalculate the size of the company after merging
this.cnt[a] += this.cnt[b];
}
}// Driver codefunction main() {
// n = number of employees & q = number of queries
const n = 3, q = 3;
const uf = new GFG(n);
// Calculate the size of node 1
console.log(uf.totalNodesInSet(1));
// Join components in which 2 and 3 are present
uf.union(2, 3);
// Calculate the size of node 3
console.log(uf.totalNodesInSet(3));
}main(); |
Output
1 2
Time Complexity: O(QlogN), N is no of employees, Q is no of queries
Auxiliary Space: O(N), N is no of employees