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

• Last Updated : 11 Feb, 2022

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][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` `        ``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)

My Personal Notes arrow_drop_up
Recommended Articles
Page :