Related Articles

# Find the array to which each element in given queries belong along with count of elements

• Last Updated : 07 Oct, 2021

Given an array of pairs arr[][] of length N, and an array queries[] of length M, and an integer R, where each query contains an integer from 1 to R, the task for every queries[i] is to find the set to which it belongs and find the total number of elements of the set.

Note: Initially every integer from 1 to R belongs to the distinct set.

Examples:

Input: R = 5, arr[] = {{1, 2}, {2, 3}, {4, 5}}, queries[] = {2, 4, 1, 3}
Output: 3 2 3 3
Explanation:
After making the sets from the arr[] pairs, {1, 2, 3}, {4, 5}
For the first query: 2 belongs to the set {1, 2, 3} and the total number of elements is 3.
For the second query: 4 belongs to the set {4, 5} and the total number of elements is 2.
For the third query: 1 belongs to the set {1, 2, 3} and the total number of elements is 3.
For the fourth query: 3 belongs to the set {1, 2, 3} and the total number of elements is 3.

Input: R = 6, arr[] = {{1, 3}, {2, 4}}, queries[] = {2, 5, 6, 1}
Output: 2 1 1 2

Approach: The given problem can be solved using the Disjoint Set Union. Initially, all the elements are in different sets, process the arr[] and do union operation on the given pairs and in union update, the total[] value for the parent element. For each query do find operation and for the returned parent element find the size of the current set as the value of total[parent]. Follow the steps below to solve the problem:

• Initialize the vectors parent(R + 1), rank(R + 1, 0), total(R + 1, 1).
• Iterate over the range [1, R+1) using the variable i and set the value of parent[I] as I.
• Iterate over the range [1, N-1] using the variable i and perform the Union Operation as Union(parent, rank, total, arr[I].first, arr[I].second).
• Iterate over the range [1, M – 1] using the variable i and perform the following steps:
• Call for function for finding the parent of the current element queries[i] as Find(parent, queries[I]).
• Print the value of total[i] as the size of the current set.

Below is the implementation of the above approach:

## C++

 `// C++ program for the above approach``#include ``using` `namespace` `std;` `// Function to perform the find operation``// of disjoint set union``int` `Find(vector<``int``>& parent, ``int` `a)``{``    ``return` `parent[a]``           ``= (parent[a] == a)``                 ``? a``                 ``: (Find(parent, parent[a]));``}` `// Function to find the Union operation``// of disjoint set union``void` `Union(vector<``int``>& parent,``           ``vector<``int``>& rank,``           ``vector<``int``>& total,``           ``int` `a, ``int` `b)``{``    ``// Find the parent of a and b``    ``a = Find(parent, a);``    ``b = Find(parent, b);` `    ``if` `(a == b)``        ``return``;` `    ``// If the rank are the same``    ``if` `(rank[a] == rank[b]) {``        ``rank[a]++;``    ``}` `    ``if` `(rank[a] < rank[b]) {``        ``int` `temp = a;``        ``a = b;``        ``b = temp;``    ``}` `    ``// Update the parent for node b``    ``parent[b] = a;` `    ``// Update the total number of``    ``// elements of a``    ``total[a] += total[b];``}` `// Function to find the total element``// of the set which belongs to the``// element queries[i]``void` `findTotNumOfSet(vector >& arr,``                     ``vector<``int``>& queries,``                     ``int` `R, ``int` `N, ``int` `M)``{` `    ``// Stores the parent elements``    ``// of the sets``    ``vector<``int``> parent(R + 1);` `    ``// Stores the rank of the sets``    ``vector<``int``> rank(R + 1, 0);` `    ``// Stores the total number of``    ``// elements of the sets``    ``vector<``int``> total(R + 1, 1);` `    ``for` `(``int` `i = 1; i < R + 1; i++) {` `        ``// Update parent[i] to i``        ``parent[i] = i;``    ``}` `    ``for` `(``int` `i = 0; i < N; i++) {` `        ``// Add the arr[i].first and``        ``// arr[i].second elements to``        ``// the same set``        ``Union(parent, rank, total,``              ``arr[i].first,``              ``arr[i].second);``    ``}` `    ``for` `(``int` `i = 0; i < M; i++) {` `        ``// Find the parent element of``        ``// the element queries[i]``        ``int` `P = Find(parent, queries[i]);` `        ``// Print the total elements of``        ``// the set which belongs to P``        ``cout << total[P] << ``" "``;``    ``}``}` `// Driver Code``int` `main()``{``    ``int` `R = 5;``    ``vector > arr{ { 1, 2 },``                                 ``{ 2, 3 },``                                 ``{ 4, 5 } };``    ``vector<``int``> queries{ 2, 4, 1, 3 };``    ``int` `N = arr.size();``    ``int` `M = queries.size();` `    ``findTotNumOfSet(arr, queries, R, N, M);` `    ``return` `0;``}`

## Java

 `// Java program for the above approach``class` `GFG``{` `// Function to perform the find operation``// of disjoint set union``static` `int` `Find(``int` `[] parent, ``int` `a)``{``    ``return` `parent[a]``           ``= (parent[a] == a)``                 ``? a``                 ``: (Find(parent, parent[a]));``}` `// Function to find the Union operation``// of disjoint set union``static` `void` `Union(``int` `[] parent,``           ``int` `[] rank,``           ``int` `[] total,``           ``int` `a, ``int` `b)``{``    ``// Find the parent of a and b``    ``a = Find(parent, a);``    ``b = Find(parent, b);` `    ``if` `(a == b)``        ``return``;` `    ``// If the rank are the same``    ``if` `(rank[a] == rank[b]) {``        ``rank[a]++;``    ``}` `    ``if` `(rank[a] < rank[b]) {``        ``int` `temp = a;``        ``a = b;``        ``b = temp;``    ``}` `    ``// Update the parent for node b``    ``parent[b] = a;` `    ``// Update the total number of``    ``// elements of a``    ``total[a] += total[b];``}` `// Function to find the total element``// of the set which belongs to the``// element queries[i]``static` `void` `findTotNumOfSet(``int``[][] arr,``                     ``int` `[] queries,``                     ``int` `R, ``int` `N, ``int` `M)``{` `    ``// Stores the parent elements``    ``// of the sets``    ``int` `[] parent = ``new` `int``[R + ``1``];` `    ``// Stores the rank of the sets``    ``int` `[] rank = ``new` `int``[R + ``1``];` `    ``// Stores the total number of``    ``// elements of the sets``    ``int` `[] total = ``new` `int``[R + ``1``];``    ``for` `(``int` `i = ``0``; i < total.length; i++) {``        ``total[i] = ``1``;``    ``}``    ``for` `(``int` `i = ``1``; i < R + ``1``; i++) {` `        ``// Update parent[i] to i``        ``parent[i] = i;``    ``}` `    ``for` `(``int` `i = ``0``; i < N; i++) {` `        ``// Add the arr[i] and``        ``// arr[i] elements to``        ``// the same set``        ``Union(parent, rank, total,``              ``arr[i][``0``],``              ``arr[i][``1``]);``    ``}` `    ``for` `(``int` `i = ``0``; i < M; i++) {` `        ``// Find the parent element of``        ``// the element queries[i]``        ``int` `P = Find(parent, queries[i]);` `        ``// Print the total elements of``        ``// the set which belongs to P``        ``System.out.print(total[P]+ ``" "``);``    ``}``}` `// Driver Code``public` `static` `void` `main(String[] args)``{``    ``int` `R = ``5``;``    ``int``[][]  arr = { { ``1``, ``2` `},``                                 ``{ ``2``, ``3` `},``                                 ``{ ``4``, ``5` `} };``    ``int` `[] queries = { ``2``, ``4``, ``1``, ``3` `};``    ``int` `N = arr.length;``    ``int` `M = queries.length;` `    ``findTotNumOfSet(arr, queries, R, N, M);` `}``}` `// This code is contributed by 29AjayKumar.`

## Python3

 `# Python3 program for the above approach` `# Function to perform the find operation``# of disjoint set union``def` `Find(parent, a):``    ``if` `(parent[a] ``=``=` `a):``        ``return` `a``    ``else``:``        ``return` `Find(parent, parent[a])` `# Function to find the Union operation``# of disjoint set union``def` `Union(parent, rank, total, a, b):``  ` `    ``# Find the parent of a and b``    ``a ``=` `Find(parent, a)``    ``b ``=` `Find(parent, b)``    ``if``(a ``=``=` `b):``        ``return``      ` `    ``# If the rank are the same``    ``if``(rank[a] ``=``=` `rank[b]):``        ``rank[a] ``+``=` `1``    ``if``(rank[a] < rank[b]):``        ``temp ``=` `a``        ``a ``=` `b``        ``b ``=` `temp``        ` `    ``# Update the parent for node b``    ``parent[b] ``=` `a``    ` `    ``# Update the total number of``    ``# elements of a``    ``total[a] ``+``=` `total[b]` `# Function to find the total element``# of the set which belongs to the``# element queries[i]``def` `findTotNumOfSet(arr, queries, R, N, M):``  ` `    ``# Stores the parent elements``    ``# of the sets``    ``parent ``=` `[``None``]``*``(R``+``1``)``    ` `    ``# Stores the rank of the sets``    ``rank ``=` `[``0``]``*``(R``+``1``)``    ` `    ``# Stores the total number of``    ``# elements of the sets``    ``total ``=` `[``1``]``*``(R ``+` `1``)``    ``for` `i ``in` `range``(``1``, R ``+` `1``):``      ` `        ``# Add the arr[i].first and``        ``# arr[i].second elements to``        ``# the same set``        ``parent[i] ``=` `i``    ``for` `i ``in` `range``(N):``        ``Union(parent, rank, total, arr[i][``0``], arr[i][``1``])``    ``for` `i ``in` `range``(M):``      ` `        ``# Find the parent element of``        ``# the element queries[i]``        ``P ``=` `Find(parent, queries[i])``        ` `        ``# Print the total elements of``        ``# the set which belongs to P``        ``print``(total[P], end``=``" "``)` `# Driver code``R ``=` `5``arr ``=` `[[``1``, ``2``], [``2``, ``3``], [``4``, ``5``]]``queries ``=` `[``2``, ``4``, ``1``, ``3``]``N ``=` `len``(arr)``M ``=` `len``(queries)``findTotNumOfSet(arr, queries, R, N, M)` `# This code is contributed by parthmanchanda81`

## C#

 `// C# program for the above approach``using` `System;` `public` `class` `GFG``{` `// Function to perform the find operation``// of disjoint set union``static` `int` `Find(``int` `[] parent, ``int` `a)``{``    ``return` `parent[a]``           ``= (parent[a] == a)``                 ``? a``                 ``: (Find(parent, parent[a]));``}` `// Function to find the Union operation``// of disjoint set union``static` `void` `Union(``int` `[] parent,``           ``int` `[] rank,``           ``int` `[] total,``           ``int` `a, ``int` `b)``{``    ``// Find the parent of a and b``    ``a = Find(parent, a);``    ``b = Find(parent, b);` `    ``if` `(a == b)``        ``return``;` `    ``// If the rank are the same``    ``if` `(rank[a] == rank[b]) {``        ``rank[a]++;``    ``}` `    ``if` `(rank[a] < rank[b]) {``        ``int` `temp = a;``        ``a = b;``        ``b = temp;``    ``}` `    ``// Update the parent for node b``    ``parent[b] = a;` `    ``// Update the total number of``    ``// elements of a``    ``total[a] += total[b];``}` `// Function to find the total element``// of the set which belongs to the``// element queries[i]``static` `void` `findTotNumOfSet(``int``[,] arr,``                     ``int` `[] queries,``                     ``int` `R, ``int` `N, ``int` `M)``{` `    ``// Stores the parent elements``    ``// of the sets``    ``int` `[] parent = ``new` `int``[R + 1];` `    ``// Stores the rank of the sets``    ``int` `[] rank = ``new` `int``[R + 1];` `    ``// Stores the total number of``    ``// elements of the sets``    ``int` `[] total = ``new` `int``[R + 1];``    ``for` `(``int` `i = 0; i < total.Length; i++) {``        ``total[i] = 1;``    ``}``    ``for` `(``int` `i = 1; i < R + 1; i++) {` `        ``// Update parent[i] to i``        ``parent[i] = i;``    ``}` `    ``for` `(``int` `i = 0; i < N; i++) {` `        ``// Add the arr[i,0] and``        ``// arr[i,1] elements to``        ``// the same set``        ``Union(parent, rank, total,``              ``arr[i,0],``              ``arr[i,1]);``    ``}` `    ``for` `(``int` `i = 0; i < M; i++) {` `        ``// Find the parent element of``        ``// the element queries[i]``        ``int` `P = Find(parent, queries[i]);` `        ``// Print the total elements of``        ``// the set which belongs to P``        ``Console.Write(total[P]+ ``" "``);``    ``}``}` `// Driver Code``public` `static` `void` `Main(String[] args)``{``    ``int` `R = 5;``    ``int``[,]  arr = { { 1, 2 },``                                 ``{ 2, 3 },``                                 ``{ 4, 5 } };``    ``int` `[] queries = { 2, 4, 1, 3 };``    ``int` `N = arr.GetLength(0);``    ``int` `M = queries.GetLength(0);` `    ``findTotNumOfSet(arr, queries, R, N, M);``}``}` `// This code is contributed by shikhasingrajput`

## Javascript

 ``
Output:
`3 2 3 3`

Time Complexity: O(M*log R)
Auxiliary Space: O(R)

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with experts, please refer DSA Live Classes for Working Professionals and Competitive Programming Live for Students.

My Personal Notes arrow_drop_up