Path queries through restricted weighted edges

Given a graph G with N vertices numbered from 1 to N and M undirected edges defined from vertex u to vertex v i.e. [ui, v_i]. Each edge has an integer weight [w_i] associated with it. You are given Q queries where each query consists of a vertex and an integer weight i.e. [a_i, z_i]. For each query print the number of vertices it can reach from a_i avoiding edges with weights less than equal to z_i.

Examples:

Input: N = 4, M = 3, G = {{1, 2, 1}, {1, 3, 2}, {3, 4, 3}}, Q = 3, Q_arr = {{1, 1}, {2, 0}, {3, 1}}
Output: {3, 4, 3}
Explanation:

• Query 1: V = 1, vertices reachable from 1 without traversing through edges with weights <= 1 are 1->3->4. Hence total 3 vertices
• Query 2: V = 3, vertices reachable from 2 without traversing through edges with weights <= 0 are 2->1->3->4. Hence total 4 vertices
• Query 3: V = 3, vertices reachable from 3 without traversing through edges with weights <= 1 are 3->4 and 3->1. Hence total 3 vertices

Input: N = 5, M = 4, G = { {1, 2, 3}, {2, 3, 5}, {3, 4, 2}, {4, 5, 4}}, Q = 3, Q_arr = {{1, 3}, {1, 2}, {3, 1}}
Output: {1, 3, 5}

Naive Approach: To solve the problem using BFS/DFS follow the below idea:

For each query, we have to in short find the size of the connected component in which the vertex a_i lie in. The connected component defined will be the set of vertices reachable from a_i such that the weights are strictly greater than z_i.

Follow the steps to solve the problem:

• Store the graph in an adjacency list using a struct or a vector along with its weights.
• Traverse through the queries.
• For each query,
  • Perform a BFS/DFS call.
  • Traverse only the adjacent vertices where the weight w_i is strictly greater than z_i.
  • Maintain count of such vertices.
• Return the count.

Below is the implementation for the above approach:

C++

// C++ code for the above approach:
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;

struct Edge {
    int u;
    int v;
    Edge(int c, int y)
        : u(c)
        , v(y)
    {
    }
};

// Adjacency list representation
// of the graph
vector<vector<Edge> > G;

// Keeps track of visited vertices
vector<bool> used;

int countComponentSize(int v, int w)
{

    // If vertex v is already visited,
    // return 0
    if (used[v]) {
        return 0;
    }

    // Start with count = 1 (for the
    // current vertex)
    int count = 1;

    // Mark the current vertex
    // as visited
    used[v] = true;

    // Iterate over the neighboring
    // vertices
    for (const Edge& e : G[v]) {
        int next = e.u;
        int year = e.v;

        // Skip edges with a weight less
        // than or equal to w
        if (year <= w) {
            continue;
        }

        // Recursively count vertices in the
        // connected component
        count += countComponentSize(next, w);
    }

    // Return the total count of vertices in
    // the connected component
    return count;
}

// Drivers code
int main()
{

    // Read the number of vertices and edges
    int n, m;
    n = 4;
    m = 3;

    // Resize the G vector to accommodate
    // n vertices
    G.resize(n + 1);

    // Initialize the used vector with
    // false for all vertices
    used.resize(n + 1, false);

    vector<vector<int> > graph
        = { { 1, 2, 1 }, { 1, 3, 2 }, { 3, 4, 3 } };

    for (int i = 0; i < graph.size(); i++) {

        // Read the edge information (vertex a,
        // vertex b, weight)
        int a = graph[i][0], b = graph[i][1],
            y = graph[i][2];

        // Add edge (b, y) to the adjacency
        // list of vertex a
        G[a].emplace_back(b, y);

        // Add edge (a, y) to the adjacency
        // list of vertex b
        G[b].emplace_back(a, y);
    }

    // Read the number of queries
    int q = 3;

    vector<vector<int> > queries
        = { { 1, 1 }, { 2, 0 }, { 3, 1 } };

    for (int i = 0; i < queries.size(); i++) {

        // Read the query (vertex v, vertex w)
        int v = queries[i][0], w = queries[i][1];

        // Reset the used vector for each query
        fill(used.begin(), used.end(), false);

        // Compute the count of vertices in
        // the connected component
        int count = countComponentSize(v, w);

        // Compute the count of vertices in
        // the connected component
        cout << count << endl;
    }

    return 0;
}

Java

// Java code for the above approach:
import java.util.ArrayList;
import java.util.List;

class Edge {
    int u;
    int v;

    public Edge(int c, int y) {
        u = c;
        v = y;
    }
}

public class ConnectedComponents {
    // Adjacency list representation of the graph
    List<List<Edge>> G;
    // Keeps track of visited vertices
    boolean[] used;

    public int countComponentSize(int v, int w) {
        // If vertex v is already visited, return 0
        if (used[v]) {
            return 0;
        }

        // Start with count = 1 (for the current vertex)
        int count = 1;

        // Mark the current vertex as visited
        used[v] = true;

        // Iterate over the neighboring vertices
        for (Edge e : G.get(v)) {
            int next = e.u;
            int year = e.v;

            // Skip edges with a weight less than or equal to w
            if (year <= w) {
                continue;
            }

            // Recursively count vertices in the connected component
            count += countComponentSize(next, w);
        }

        // Return the total count of vertices in the connected component
        return count;
    }

    public static void main(String[] args) {
        // Read the number of vertices and edges
        int n = 4;
        int m = 3;

        // Resize the G list to accommodate n vertices
        List<List<Edge>> G = new ArrayList<>();
        for (int i = 0; i <= n; i++) {
            G.add(new ArrayList<>());
        }

        // Initialize the used array with false for all vertices
        boolean[] used = new boolean[n + 1];

        List<int[]> graph =
          List.of(new int[]{1, 2, 1}, new int[]{1, 3, 2}, new int[]{3, 4, 3});

        for (int i = 0; i < graph.size(); i++) {
            // Read the edge information (vertex a, vertex b, weight)
            int a = graph.get(i)[0], b = graph.get(i)[1], y = graph.get(i)[2];

            // Add edge (b, y) to the adjacency list of vertex a
            G.get(a).add(new Edge(b, y));

            // Add edge (a, y) to the adjacency list of vertex b
            G.get(b).add(new Edge(a, y));
        }

        // Read the number of queries
        int q = 3;

        List<int[]> queries =
          List.of(new int[]{1, 1}, new int[]{2, 0}, new int[]{3, 1});

        for (int i = 0; i < queries.size(); i++) {
            // Read the query (vertex v, vertex w)
            int v = queries.get(i)[0], w = queries.get(i)[1];

            // Reset the used array for each query
            for (int j = 0; j <= n; j++) {
                used[j] = false;
            }

            // Compute the count of vertices in the connected component
            ConnectedComponents cc = new ConnectedComponents();
            cc.G = G;
            cc.used = used;
            int count = cc.countComponentSize(v, w);

            // Compute the count of vertices in the connected component
            System.out.println(count);
        }
    }
}

Python3

class Edge:
    def __init__(self, u, v):
        self.u = u
        self.v = v


# Adjacency list representation of the graph
G = []

# Keeps track of visited vertices
used = []


def countComponentSize(v, w):
    # If vertex v is already visited, return 0
    if used[v]:
        return 0

    # Start with count = 1 (for the current vertex)
    count = 1

    # Mark the current vertex as visited
    used[v] = True

    # Iterate over the neighboring vertices
    for e in G[v]:
        next = e.u
        year = e.v

        # Skip edges with a weight less than or equal to w
        if year <= w:
            continue

        # Recursively count vertices in the connected component
        count += countComponentSize(next, w)

    # Return the total count of vertices in the connected component
    return count


if __name__ == "__main__":
    # Read the number of vertices and edges
    n, m = 4, 3

    # Create adjacency list representation of the graph
    G = [[] for _ in range(n + 1)]

    # Initialize the used list with False for all vertices
    used = [False] * (n + 1)

    # Define the edges of the graph
    graph = [[1, 2, 1], [1, 3, 2], [3, 4, 3]]

    # Add the edges to the adjacency list
    for a, b, y in graph:
        # Add edge (b, y) to the adjacency list of vertex a
        G[a].append(Edge(b, y))

        # Add edge (a, y) to the adjacency list of vertex b
        G[b].append(Edge(a, y))

    # Read the number of queries
    q = 3

    # Define the queries
    queries = [[1, 1], [2, 0], [3, 1]]

    # Process each query
    for v, w in queries:
        # Reset the used list for each query
        used = [False] * (n + 1)

        # Compute the count of vertices in the connected component
        count = countComponentSize(v, w)

        # Print the count of vertices in the connected component
        print(count)

C#

using System;
using System.Collections.Generic;

class Edge
{
    public int U { get; set; }
    public int V { get; set; }

    public Edge(int u, int v)
    {
        U = u;
        V = v;
    }
}

public class ConnectedComponents
{
    // Adjacency list representation of the graph
    private List<List<Edge>> G;
    // Keeps track of visited vertices
    private bool[] Used;

    public int CountComponentSize(int v, int w)
    {
        // If vertex v is already visited, return 0
        if (Used[v])
        {
            return 0;
        }

        // Start with count = 1 (for the current vertex)
        int count = 1;

        // Mark the current vertex as visited
        Used[v] = true;

        // Iterate over the neighboring vertices
        foreach (Edge e in G[v])
        {
            int next = e.U;
            int year = e.V;

            // Skip edges with a weight less than or equal to w
            if (year <= w)
            {
                continue;
            }

            // Recursively count vertices in the connected component
            count += CountComponentSize(next, w);
        }

        // Return the total count of vertices in the connected component
        return count;
    }

    public static void Main(string[] args)
    {
        // Read the number of vertices and edges
        int n = 4;

        // Resize the G list to accommodate n vertices
        List<List<Edge>> G = new List<List<Edge>>();
        for (int i = 0; i <= n; i++)
        {
            G.Add(new List<Edge>());
        }

        // Initialize the Used array with false for all vertices
        bool[] Used = new bool[n + 1];

        List<int[]> graph =
            new List<int[]>
            {
                new int[] { 1, 2, 1 },
                new int[] { 1, 3, 2 },
                new int[] { 3, 4, 3 }
            };

        for (int i = 0; i < graph.Count; i++)
        {
            // Read the edge information (vertex a, vertex b, weight)
            int a = graph[i][0], b = graph[i][1], y = graph[i][2];

            // Add edge (b, y) to the adjacency list of vertex a
            G[a].Add(new Edge(b, y));

            // Add edge (a, y) to the adjacency list of vertex b
            G[b].Add(new Edge(a, y));
        }

        // Read the number of queries

        List<int[]> queries =
            new List<int[]>
            {
                new int[] { 1, 1 },
                new int[] { 2, 0 },
                new int[] { 3, 1 }
            };

        for (int i = 0; i < queries.Count; i++)
        {
            // Read the query (vertex v, vertex w)
            int v = queries[i][0], w = queries[i][1];

            // Reset the Used array for each query
            for (int j = 0; j <= n; j++)
            {
                Used[j] = false;
            }

            // Compute the count of vertices in the connected component
            ConnectedComponents cc = new ConnectedComponents();
            cc.G = G;
            cc.Used = Used;
            int count = cc.CountComponentSize(v, w);

            // Compute the count of vertices in the connected component
            Console.WriteLine(count);
        }
    }
}

Javascript

// JavaScript code for above approach

// Define the Edge class
class Edge {
    constructor(u, v) {
        this.u = u;
        this.v = v;
    }
}

// Adjacency list representation
// of the graph
const G = [];

// Keeps track of visited vertices
const used = [];

// Function to count the size of
// the connected component
function countComponentSize(v, w) {
    // If vertex v is already
    // visited, return 0
    if (used[v]) {
        return 0;
    }

    // Start with count = 1 (for
    // the current vertex)
    let count = 1;

    // Mark the current vertex as visited
    used[v] = true;

    // Iterate over the neighboring vertices
    for (const e of G[v]) {
        const next = e.u;
        const year = e.v;

        // Skip edges with a weight
        // less than or equal to w
        if (year <= w) {
            continue;
        }

        // Recursively count vertices in
        // the connected component
        count += countComponentSize(next, w);
    }

    // Return the total count of vertices
    // in the connected component
    return count;
}

// Driver code

// Read the number of vertices and edges
const n = 4;
const m = 3;

// Resize the G vector to accommodate
// n vertices
for (let i = 0; i <= n; i++) {
    G.push([]);
}

// Initialize the used vector with
// false for all vertices
for (let i = 0; i <= n; i++) {
    used.push(false);
}

const graph = [[1, 2, 1], [1, 3, 2], [3, 4, 3]];

for (let i = 0; i < graph.length; i++) {
    // Read the edge information
    // (vertex a, vertex b, weight)
    const a = graph[i][0];
    const b = graph[i][1];
    const y = graph[i][2];

    // Add edge (b, y) to the
    // adjacency list of vertex a
    G[a].push(new Edge(b, y));

    // Add edge (a, y) to the
    // adjacency list of vertex b
    G[b].push(new Edge(a, y));
}

// Read the number of queries
const q = 3;

const queries = [[1, 1], [2, 0], [3, 1]];

for (let i = 0; i < queries.length; i++) {
    // Read the query (vertex v, vertex w)
    const v = queries[i][0];
    const w = queries[i][1];

    // Reset the used vector for each query
    used.fill(false);

    // Compute the count of vertices
    // in the connected component
    const count = countComponentSize(v, w);

    // Compute the count of vertices
    // in the connected component
    console.log(count);
}

// This code is contributed by prasad264

Output

3
4
3

Time Complexity: O(Q*(N+M))
Auxiliary Space: O(M)

Efficient Approach: To solve the problem using Disjoint Set follow the below idea:

For each qth query, maintain a union find set with each individual set having vertices whose weights are strictly greater than z_i. Sort the queries and graph in the decreasing order of weights since union-find can only merge the two connected components but not disconnect them. Hence with each query sorted in decreasing order the union operation needs to be performed with the previously solved query since the weights will decrease and not increase.

Follow the steps to solve the problem:

• Define the UnionFind class to represent a disjoint set data structure.
• Implement the constructor of the UnionFind class to initialize the parent and size arrays.
• Implement the findParent method in the UnionFind class to find the representative (root) of a set.
• Implement the mergeSets method in the UnionFind class to merge two sets.
• Implement the getSize method in the UnionFind class to get the size of a set.
• Define the Edge struct to represent an edge in the graph.
• Define the Query struct to represent a query.
• In the main function, declare, and initialize the necessary variables such as the number of vertices, edges, and queries.
• Create a processed_edges array to store the query and graph information. While storing the query information, make sure to insert the order of query to access the results in order as that of query asked.
• Sort the processed edges in descending order of weight; if weights are equal, queries should come before edges.
• Create an instance of the UnionFind class with the given number of vertices.
• Process each object of processed_edge:
  • If its a query then,
    • Extract the vertex.
    • Find the size of the connected component that the vertex lies in using the getSize method of the UnionFind class.
    • Store the size at the index extracted from the object into the resultant array at that index.
  • Else is a graph construction,
    • Find the two vertices U and V
    • Merge them using the merge function
• Output the sizes of the connected components in order.

C++

// C++ code for the above approach:
#include <bits/stdc++.h>
using namespace std;

class UnionFind {
    vector<int> parent, size;

public:
    void initialize(int n)
    {

        // Create a vector to store parent
        // information for each node
        parent.resize(n + 2);

        // Create a vector to store size
        // information for each node
        // (initialized to 1)
        size.resize(n + 2, 1);
        for (int i = 0; i <= n; i++) {

            // Initially, set the parent of
            // each node as itself
            parent[i] = i;
        }
    }

    int find(int u)
    {
        if (u == parent[u]) {

            // If the current node is the
            // parent itself, return the node
            return u;
        }

        // Otherwise, recursively find the
        // parent of the current node and
        // update the parent of the current
        // node
        return parent[u] = find(parent[u]);
    }

    void merge(int u, int v)
    {

        // Find the parent of node u
        u = find(u);

        // Find the parent of node v
        v = find(v);
        if (u == v) {

            // If both nodes have the same
            // parent, they are already in
            // the same set, so no need to merge
            return;
        }
        else {

            // Update the size of the set
            // of node u by adding the size
            // of the set of node v
            size[u] += size[v];

            // Make the parent of node v as u,
            // indicating that they are in
            // the same set
            parent[v] = u;
        }
    }

    int getSize(int X)
    {

        // Return the size of the set to
        // which node X belongs
        return size[X];
    }
};

vector<vector<int> >
process_graph(vector<vector<int> > G,
              vector<vector<int> > query, int N, int M)
{

    // Create a 2D vector to store
    // the processed edges
    vector<vector<int> > processed_edges;
    vector<int> temp;

    for (auto edge : G) {
        int u = edge[0], v = edge[1], w = edge[2];

        // Add the edge to the processed
        // edges vector with a flag '0'
        // indicating it is an edge from the
        // original graph
        processed_edges.push_back({ w, u, v, 0 });
    }

    int index = 0;
    for (auto q : query) {
        int a = q[0];
        int z = q[1];

        // Add the query to the processed
        // edges vector with a flag '1'
        // indicating it is a query
        processed_edges.push_back({ z, index, a, 1 });
        index++;
    }

    // Sort the processed edges in descending
    // order of weight, and if weights are
    // equal, queries should come
    // before edges
    sort(processed_edges.begin(), processed_edges.end(),
         [&](vector<int> a, vector<int> b) {
             return a[0] > b[0]
                    || (a[0] == b[0] and a[3] > b[3]);
         });

    // Return the processed edges vector
    return processed_edges;
}

void solve(int N, int M, vector<vector<int> > G,
           vector<vector<int> > query)
{

    // Process the graph and queries
    vector<vector<int> > graph_sorted
        = process_graph(G, query, N, M);

    UnionFind UF;

    // Initialize the UnionFind data
    // structure with N nodes
    UF.initialize(N);

    // Create a vector to store the
    // results of the queries
    vector<int> res(query.size());

    for (auto itr : graph_sorted) {
        if (itr[3] == 1) {
            int a = itr[2], w = itr[0], index = itr[1];

            // For queries, find the size of
            // the set to which node 'a' belongs
            // and store it in the result vector
            res[index] = UF.getSize(UF.find(a));
        }
        else {
            int u = itr[1], v = itr[2];

            // For edges, merge the sets to
            // which nodes 'u' and 'v' belong
            UF.merge(u, v);
        }
    }

    for (auto answer : res) {

        // Print the results of the queries
        cout << answer << "\n";
    }
}

// Drivers code
int main()
{
    int N = 4, M = 3;
    vector<vector<int> > G
        = { { 1, 2, 1 }, { 1, 3, 2 }, { 3, 4, 3 } };
    int Q = 3;
    vector<vector<int> > query
        = { { 1, 1 }, { 2, 0 }, { 3, 1 } };

    // Function Call
    solve(N, M, G, query);

    return 0;
}

Java

// Java code for the above approach
import java.util.*;

class UnionFind {
    private List<Integer> parent, size;

    public void initialize(int n)
    {
        // Create a list to store parent information for
        // each node
        parent = new ArrayList<>(
            Collections.nCopies(n + 2, 0));

        // Create a list to store size information for each
        // node (initialized to 1)
        size = new ArrayList<>(
            Collections.nCopies(n + 2, 1));

        for (int i = 0; i <= n; i++) {
            // Initially, set the parent of each node as
            // itself
            parent.set(i, i);
        }
    }

    public int find(int u)
    {
        if (u == parent.get(u)) {
            // If the current node is the parent itself,
            // return the node
            return u;
        }

        // Otherwise, recursively find the parent of the
        // current node and update the parent of the current
        // node
        parent.set(u, find(parent.get(u)));
        return parent.get(u);
    }

    public void merge(int u, int v)
    {
        // Find the parent of node u
        u = find(u);

        // Find the parent of node v
        v = find(v);
        if (u == v) {
            // If both nodes have the same parent, they are
            // already in the same set, so no need to merge
            return;
        }

        // Update the size of the set of node u by adding
        // the size of the set of node v
        size.set(u, size.get(u) + size.get(v));

        // Make the parent of node v as u, indicating that
        // they are in the same set
        parent.set(v, u);
    }

    public int getSize(int X)
    {
        // Return the size of the set to which node X
        // belongs
        return size.get(find(X));
    }
}

public class GFG {
    public static List<List<Integer> >
    processGraph(List<List<Integer> > G,
                 List<List<Integer> > query, int N, int M)
    {
        // Create a 2D list to store the processed edges
        List<List<Integer> > processedEdges
            = new ArrayList<>();

        for (List<Integer> edge : G) {
            int u = edge.get(0), v = edge.get(1),
                w = edge.get(2);

            // Add the edge to the processed edges list with
            // a flag '0' indicating it is an edge from the
            // original graph
            processedEdges.add(Arrays.asList(w, u, v, 0));
        }

        int index = 0;
        for (List<Integer> q : query) {
            int a = q.get(0);
            int z = q.get(1);

            // Add the query to the processed edges list
            // with a flag '1' indicating it is a query
            processedEdges.add(
                Arrays.asList(z, index, a, 1));
            index++;
        }

        // Sort the processed edges in descending order of
        // weight, and if weights are equal, queries should
        // come before edges
        processedEdges.sort((a, b) -> {
            if (a.get(0).equals(b.get(0))) {
                return Integer.compare(b.get(3), a.get(3));
            }
            return Integer.compare(b.get(0), a.get(0));
        });

        // Return the processed edges list
        return processedEdges;
    }

    public static void solve(int N, int M,
                             List<List<Integer> > G,
                             List<List<Integer> > query)
    {
        // Process the graph and queries
        List<List<Integer> > graphSorted
            = processGraph(G, query, N, M);
        UnionFind UF = new UnionFind();

        // Initialize the UnionFind data structure with N
        // nodes
        UF.initialize(N);

        // Create a list to store the results of the queries
        List<Integer> res = new ArrayList<>(
            Collections.nCopies(query.size(), 0));

        for (List<Integer> itr : graphSorted) {
            if (itr.get(3) == 1) {
                int a = itr.get(2), w = itr.get(0),
                    index = itr.get(1);

                // For queries, find the size of the set to
                // which node 'a' belongs and store it in
                // the result list
                res.set(index, UF.getSize(UF.find(a)));
            }
            else {
                int u = itr.get(1), v = itr.get(2);

                // For edges, merge the sets to which nodes
                // 'u' and 'v' belong
                UF.merge(u, v);
            }
        }

        for (int answer : res) {
            // Print the results of the queries
            System.out.println(answer);
        }
    }

    // Driver code
    public static void main(String[] args)
    {
        int N = 4, M = 3;
        List<List<Integer> > G = Arrays.asList(
            Arrays.asList(1, 2, 1), Arrays.asList(1, 3, 2),
            Arrays.asList(3, 4, 3));
        List<List<Integer> > query = Arrays.asList(
            Arrays.asList(1, 1), Arrays.asList(2, 0),
            Arrays.asList(3, 1));

        // Function Call
        solve(N, M, G, query);
    }
}

// This code is contributed by Susobhan Akhuli

Python3

class UnionFind:
    def initialize(self, n):
        # Create a list to store parent information for each node
        self.parent = list(range(n + 2))
        # Create a list to store size information for each node (initialized to 1)
        self.size = [1] * (n + 2)

    def find(self, u):
        if u == self.parent[u]:
            # If the current node is the parent itself, return the node
            return u
        # Otherwise, recursively find the parent of the current node and update the parent of the current node
        self.parent[u] = self.find(self.parent[u])
        return self.parent[u]

    def merge(self, u, v):
        # Find the parent of node u
        u = self.find(u)
        # Find the parent of node v
        v = self.find(v)
        if u == v:
            # If both nodes have the same parent, they are already in the same set, so no need to merge
            return
        else:
            # Update the size of the set of node u by adding the size of the set of node v
            self.size[u] += self.size[v]
            # Make the parent of node v as u, indicating that they are in the same set
            self.parent[v] = u

    def get_size(self, X):
        # Return the size of the set to which node X belongs
        return self.size[X]


def process_graph(G, query, N, M):
    # Create a list to store the processed edges
    processed_edges = []
    for edge in G:
        u, v, w = edge
        # Add the edge to the processed edges list with a flag '0' indicating it is an edge from the original graph
        processed_edges.append([w, u, v, 0])

    index = 0
    for q in query:
        a, z = q
        # Add the query to the processed edges list with a flag '1' indicating it is a query
        processed_edges.append([z, index, a, 1])
        index += 1

    # Sort the processed edges in descending order of weight, and if weights are equal, queries should come before edges
    processed_edges.sort(key=lambda x: (x[0], x[3]), reverse=True)
    # Return the processed edges list
    return processed_edges


def solve(N, M, G, query):
    # Process the graph and queries
    graph_sorted = process_graph(G, query, N, M)

    UF = UnionFind()
    UF.initialize(N)  # Initialize the UnionFind data structure with N nodes

    # Create a list to store the results of the queries
    res = [0] * len(query)

    for itr in graph_sorted:
        if itr[3] == 1:
            a, w, index = itr[2], itr[0], itr[1]
            # For queries, find the size of the set to which node 'a' belongs and store it in the result list
            res[index] = UF.get_size(UF.find(a))
        else:
            u, v = itr[1], itr[2]
            # For edges, merge the sets to which nodes 'u' and 'v' belong
            UF.merge(u, v)

    for answer in res:
        # Print the results of the queries
        print(answer)


if __name__ == "__main__":
    N, M = 4, 3
    G = [[1, 2, 1], [1, 3, 2], [3, 4, 3]]
    query = [[1, 1], [2, 0], [3, 1]]

    solve(N, M, G, query)

C#

// C# code for the above approach:
using System;
using System.Collections.Generic;
using System.Linq;

class UnionFind
{
    List<int> parent, size;

    public void Initialize(int n)
    {
        // Create a list to store parent
        // information for each node
        parent = new List<int>(Enumerable.Repeat(0, n + 2));

        // Create a list to store size
        // information for each node
        // (initialized to 1)
        size = new List<int>(Enumerable.Repeat(1, n + 2));

        for (int i = 0; i <= n; i++)
        {
            // Initially, set the parent of
            // each node as itself
            parent[i] = i;
        }
    }

    public int Find(int u)
    {
        if (u == parent[u])
        {
            // If the current node is the
            // parent itself, return the node
            return u;
        }

        // Otherwise, recursively find the
        // parent of the current node and
        // update the parent of the current
        // node
        parent[u] = Find(parent[u]);
        return parent[u];
    }

    public void Merge(int u, int v)
    {
        // Find the parent of node u
        u = Find(u);

        // Find the parent of node v
        v = Find(v);
        if (u == v)
        {
            // If both nodes have the same
            // parent, they are already in
            // the same set, so no need to merge
            return;
        }

        // Update the size of the set
        // of node u by adding the size
        // of the set of node v
        size[u] += size[v];

        // Make the parent of node v as u,
        // indicating that they are in
        // the same set
        parent[v] = u;
    }

    public int GetSize(int X)
    {
        // Return the size of the set to
        // which node X belongs
        return size[Find(X)];
    }
}

public class GFG
{
    static List<List<int>> ProcessGraph(List<List<int>> G, List<List<int>> query, int N, int M)
    {
        // Create a 2D list to store
        // the processed edges
        List<List<int>> processedEdges = new List<List<int>>();

        foreach (var edge in G)
        {
            int u = edge[0], v = edge[1], w = edge[2];

            // Add the edge to the processed
            // edges list with a flag '0'
            // indicating it is an edge from the
            // original graph
            processedEdges.Add(new List<int> { w, u, v, 0 });
        }

        int index = 0;
        foreach (var q in query)
        {
            int a = q[0];
            int z = q[1];

            // Add the query to the processed
            // edges list with a flag '1'
            // indicating it is a query
            processedEdges.Add(new List<int> { z, index, a, 1 });
            index++;
        }

        // Sort the processed edges in descending
        // order of weight, and if weights are
        // equal, queries should come
        // before edges
        processedEdges.Sort((a, b) =>
        {
            if (a[0] == b[0])
            {
                return a[3] == b[3] ? 0 : (a[3] > b[3] ? -1 : 1);
            }
            return b[0] - a[0];
        });

        // Return the processed edges list
        return processedEdges;
    }

    static void Solve(int N, int M, List<List<int>> G, List<List<int>> query)
    {
        // Process the graph and queries
        List<List<int>> graphSorted = ProcessGraph(G, query, N, M);
        UnionFind UF = new UnionFind();

        // Initialize the UnionFind data
        // structure with N nodes
        UF.Initialize(N);

        // Create a list to store the
        // results of the queries
        List<int> res = new List<int>(new int[query.Count]);

        foreach (var itr in graphSorted)
        {
            if (itr[3] == 1)
            {
                int a = itr[2], w = itr[0], index = itr[1];

                // For queries, find the size of
                // the set to which node 'a' belongs
                // and store it in the result list
                res[index] = UF.GetSize(UF.Find(a));
            }
            else
            {
                int u = itr[1], v = itr[2];

                // For edges, merge the sets to
                // which nodes 'u' and 'v' belong
                UF.Merge(u, v);
            }
        }

        foreach (var answer in res)
        {
            // Print the results of the queries
            Console.WriteLine(answer);
        }
    }

    // Drivers code
    public static void Main()
    {
        int N = 4, M = 3;
        List<List<int>> G = new List<List<int>>
        {
            new List<int> { 1, 2, 1 },
            new List<int> { 1, 3, 2 },
            new List<int> { 3, 4, 3 }
        };
        List<List<int>> query = new List<List<int>>
        {
            new List<int> { 1, 1 },
            new List<int> { 2, 0 },
            new List<int> { 3, 1 }
        };

        // Function Call
        Solve(N, M, G, query);
    }
}

Javascript

class UnionFind {
  constructor(n) {
    // Create an array to store parent information for each node
    this.parent = [...Array(n + 2).keys()];
    // Create an array to store size information for each node (initialized to 1)
    this.size = new Array(n + 2).fill(1);
  }

  find(u) {
    if (u === this.parent[u]) {
      // If the current node is the parent itself, return the node
      return u;
    }
    // Otherwise, recursively find the parent of the current node and update the parent of the current node
    this.parent[u] = this.find(this.parent[u]);
    return this.parent[u];
  }

  merge(u, v) {
    // Find the parent of node u
    u = this.find(u);
    // Find the parent of node v
    v = this.find(v);
    if (u === v) {
      // If both nodes have the same parent, they are already in the same set, so no need to merge
      return;
    } else {
      // Update the size of the set of node u by adding the size of the set of node v
      this.size[u] += this.size[v];
      // Make the parent of node v as u, indicating that they are in the same set
      this.parent[v] = u;
    }
  }

  getSize(X) {
    // Return the size of the set to which node X belongs
    return this.size[X];
  }
}

function processGraph(G, query, N, M) {
  // Create an array to store the processed edges
  const processedEdges = [];
  for (const edge of G) {
    const [u, v, w] = edge;
    // Add the edge to the processed edges array with a flag '0' indicating it is an edge from the original graph
    processedEdges.push([w, u, v, 0]);
  }

  let index = 0;
  for (const q of query) {
    const [a, z] = q;
    // Add the query to the processed edges array with a flag '1' indicating it is a query
    processedEdges.push([z, index, a, 1]);
    index++;
  }

  // Sort the processed edges in descending order of weight, and if weights are equal, queries should come before edges
  processedEdges.sort((a, b) => b[0] - a[0] || a[3] - b[3]);
  // Return the processed edges array
  return processedEdges;
}

function solve(N, M, G, query) {
  // Process the graph and queries
  const graphSorted = processGraph(G, query, N, M);

  const UF = new UnionFind();
  UF.initialize(N); // Initialize the UnionFind data structure with N nodes

  // Create an array to store the results of the queries
  const res = new Array(query.length).fill(0);

  for (const itr of graphSorted) {
    if (itr[3] === 1) {
      const [a, w, index] = [itr[2], itr[0], itr[1]];
      // For queries, find the size of the set to which node 'a' belongs and store it in the result array
      res[index] = UF.getSize(UF.find(a));
    } else {
      const [u, v] = [itr[1], itr[2]];
      // For edges, merge the sets to which nodes 'u' and 'v' belong
      UF.merge(u, v);
    }
  }

  for (const answer of res) {
    // Print the results of the queries
    console.log(answer);
  }
}

if (require.main === module) {
  const N = 4;
  const M = 3;
  const G = [[1, 2, 1], [1, 3, 2], [3, 4, 3]];
  const query = [[1, 1], [2, 0], [3, 1]];

  solve(N, M, G, query);
}
// This code is contributed by shivamgupta0987654321

Output

3
4
3

Time Complexity: O(Z + ZlogZ), where Z = (Q+M).
Auxiliary Space: O(Z)