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

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 contributed by saurabh_jaiswal.
</script>

Output:

3 2 3 3

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