Skip to content
Related Articles

Related Articles

Improve Article

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

  • Last Updated : 07 Oct, 2021
Geek Week

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 <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)

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
Recommended Articles
Page :