Skip to content
Related Articles

Related Articles

Table of Contents

Improve Article
Save Article
Like Article

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.

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.

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 <bits/stdc++.h>
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<pair<int, int> >& 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<pair<int, int> > 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



<script>
// Javascript program for the above approach
 
// Function to perform the find operation
// of disjoint set union
function Find(parent, a)
{
  return (parent[a] = parent[a] == a ? a : Find(parent, parent[a]));
}
 
// Function to find the Union operation
// of disjoint set union
function 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]++;
  }
 
  if (rank[a] < rank[b]) {
    let 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]
function findTotNumOfSet(arr, queries, R, N, M)
{
 
  // Stores the parent elements
  // of the sets
  let parent = new Array(R + 1);
 
  // Stores the rank of the sets
  let rank = new Array(R + 1).fill(0);
 
  // Stores the total number of
  // elements of the sets
  let total = new Array(R + 1).fill(1);
 
  for (let i = 1; i < R + 1; i++)
  {
   
    // Update parent[i] to i
    parent[i] = i;
  }
 
  for (let 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][0], arr[i][1]);
  }
 
  for (let i = 0; i < M; i++)
  {
   
    // Find the parent element of
    // the element queries[i]
    let P = Find(parent, queries[i]);
 
    // Print the total elements of
    // the set which belongs to P
    document.write(total[P] + " ");
  }
}
 
// Driver Code
let R = 5;
let arr = [
  [1, 2],
  [2, 3],
  [4, 5],
];
let queries = [2, 4, 1, 3];
let N = arr.length;
let M = queries.length;
 
findTotNumOfSet(arr, queries, R, N, M);
 
// This code is ontributed by saurabh_jaiswal.
</script>
Output: 
3 2 3 3
 
Time Complexity: O(M*log R)
Auxiliary Space: O(R)



My Personal Notes
arrow_drop_up
Previous
Check if a Regular Bracket Sequence can be formed with concatenation of given strings
Next

Smallest vertex in the connected components of all the vertices in given undirect graph
Recommended Articles
Page :
Find the coordinate that does not belong to any square
22, Jul 21
Maximum array elements that can be removed along with its adjacent values to empty given array
05, Nov 20
Form an array of distinct elements with each element as sum of an element from each array
26, Mar 20
Queries to print count of distinct array elements after replacing element at index P by a given element
16, Oct 20
Shortest path from source to destination such that edge weights along path are alternatively increasing and decreasing
24, Apr 19
Maximize the number of uncolored vertices appearing along the path from root vertex and the colored vertices
18, Apr 20
Remaining array element after repeated removal of last element and subtraction of each element from next adjacent element
23, Feb 21
Find Array formed by adding each element of given array with largest element in new array to its left
24, Mar 20
Count of elements which are power of 2 in a given range subarray for Q queries
21, May 20
Find the element before which all the elements are smaller than it, and after which all are greater
24, Oct 15
Queries to find the maximum array element after removing elements from a given range
01, Feb 21
Find Median for each Array element by excluding the index at which Median is calculated
04, Aug 21
Sum of product of each element with each element after it
24, Mar 18
Queries to update each element in subarray to Bitwise XOR with a given value
17, Jul 20
Minimum value by which each Array element must be added as per given conditions
13, Aug 19
Print modified array after performing queries to add (i - L + 1) to each element present in the range [L, R]
19, Mar 21
Find the next non-zero Array element to the right of each array element
15, Oct 21
Find just strictly greater element from first array for each element in second array
05, May 20
Queries to find the maximum and minimum array elements excluding elements from a given range
03, Feb 21
Array formed from difference of each element from the largest element in the given array
04, Feb 20
For each element in 1st array count elements less than or equal to it in 2nd array
11, Jun 17
For each element in 1st array count elements less than or equal to it in 2nd array | Set 2
01, Oct 18
Count remaining array elements after reversing binary representation of each array element
16, Feb 21
Count of elements in first Array greater than second Array with each element considered only once
07, May 20
Article Contributed By :
https://media.geeksforgeeks.org/auth/avatar.png
dharanendralv23
@dharanendralv23
Vote for difficulty





Improved By :
Article Tags :
Practice Tags :
Report Issue

We use cookies to ensure you have the best browsing experience on our website. By using our site, you
acknowledge that you have read and understood our
Cookie Policy &
        Privacy Policy

Lightbox

Start Your Coding Journey Now!