Find the weight of the minimum spanning tree

Given a connected undirected weighted graph with N nodes and M edges. The task is to perform given queries and find the weight of the minimum spanning tree. Queries are of three types:

query(1) -> Find the weight of the minimum spanning tree.
query(2, x, y) -> Change the weight of the edge between the nodes x and y to 0.
query(3, x, y) -> Restore the weight of the edge between the nodes x and y to its original weight.

Examples:

Input:

query(2, 1, 2),
query(1),
query(3, 1, 2),
query(1)
Output:
2
3
Input:

query(1),
query(2, 3, 4),
query(1)
Output :
4
2

Approach: Let’s first compute MST of the initial graph before performing any queries and let T be this MST. The crucial observation is that at any point while handling the queries, the weight of the MST of the current graph can be computed by running Kruskal’s algorithm on edges with zero weight at this point and edges of T. So, keep edges with weight zero in a data structure and after query of type 2 and type 3 compute weight of minimum spanning tree.
Below is the implementation of the above approach:

C++

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;
#define N 2005

// To store vertices, edges
// and the required answer
int n, m, ans;

// To store parent and rank
int par[N], Rank[N];

// To store edges and the edges in MST
vector<pair<int, pair<int, int> > > edges, mst;

// To store the edges with weight zero
queue<pair<int, int> > zeros;

// Function for initialize
void initialize()
{
    for (int i = 0; i <= n; i++) {
        par[i] = i;
        Rank[i] = 0;
    }
}

// Function to add edges
void Add_edge(int u, int v, int weight)
{
    edges.push_back({ weight, { u, v } });
}

// Utility function to find set of an element i
// (uses path compression technique)
int find(int x)
{
    if (par[x] != x)
        par[x] = find(par[x]);

    return par[x];
}

// Function that performs union of two sets x and y
// (uses union by rank)
void Union(int x, int y)
{
    int xroot = find(x);
    int yroot = find(y);

    if (Rank[xroot] < Rank[yroot])
        par[xroot] = yroot;
    else if (Rank[xroot] > Rank[yroot])
        par[yroot] = xroot;
    else {
        par[yroot] = xroot;
        Rank[xroot]++;
    }
}

// Function to compute minimum spanning tree
void compute_MST()
{
    // Sort edges in increasing order of weight
    sort(edges.begin(), edges.end());

    // Go through all the edges
    for (int i = 0; i < m; i++) {
        int u = find(edges[i].second.first);
        int v = find(edges[i].second.second);

        if (u == v)
            continue;

        // Build minimum spanning tree
        // and store minimum cost
        mst.push_back(edges[i]);
        ans += edges[i].first;
        Union(u, v);
    }
}

// Function to find the cost of minimum
// spanning tree
void Modified_Kruskal(pair<int, int> x)
{
    initialize();

    // Make answer zero
    ans = 0;
    int sz = zeros.size();

    // Keep the edges which only have zero weights
    // and remove all the other edges
    for (int i = 0; i < sz; i++) {
        pair<int, int> Front = zeros.front();
        zeros.pop();

        if (Front.first == x.first
            and Front.second == x.second)
            continue;

        // Make union between the vertices of
        // edges which have weight zero and keep
        // them in queue
        Union(Front.first, Front.second);
        zeros.push(Front);
    }

    // Find the cost of the minimum spanning tree
    for (int i = 0; i < mst.size(); i++) {
        int u = find(mst[i].second.first);
        int v = find(mst[i].second.second);

        if (u == v)
            continue;

        ans += mst[i].first;
        Union(u, v);
    }
}

// Function to handle different queries
void query(int type, int u = 0, int v = 0)
{

    // Update edge weight to 0
    if (type == 2) {
        // push edge in zeros
        zeros.push({ u, v });
        Modified_Kruskal({ -1, -1 });
    }

    // Restore edge weight to original value
    else if (type == 3) {
        // push edge in zeros
        zeros.push({ u, v });
        Modified_Kruskal({ u, v });
    }
    else
        cout << ans << endl;
}

// Driver code
int main()
{

    // Number of nodes and edges
    n = 4, m = 4;
    initialize();

    // Add edges
    Add_edge(1, 2, 1);
    Add_edge(2, 3, 1);
    Add_edge(3, 4, 1);
    Add_edge(4, 1, 1);

    // Build the minimum spanning tree
    compute_MST();

    // Execute queries
    query(2, 1, 2);
    query(1);
    query(3, 1, 2);
    query(1);

    return 0;
}

Java

import java.util.*;

class Main {
    static final int N = 2005;
    static int n, m, ans;
    static int[] par, Rank;
    static List<Pair> edges, mst;
    static Queue<Pair> zeros;

    // Pair class to represent edges and vertices
    static class Pair {
        int first, second;

        Pair(int first, int second) {
            this.first = first;
            this.second = second;
        }
    }

    // Function to initialize data structures
    static void initialize() {
        par = new int[N];
        Rank = new int[N];
        for (int i = 0; i <= n; i++) {
            par[i] = i;
            Rank[i] = 0;
        }
    }

    // Function to add edges
    static void Add_edge(int u, int v, int weight) {
        edges.add(new Pair(weight, (u << 10) | v));
    }

    // Utility function to find set of an element i (uses path compression technique)
    static int find(int x) {
        if (par[x] != x)
            par[x] = find(par[x]);
        return par[x];
    }

    // Function that performs union of two sets x and y (uses union by rank)
    static void Union(int x, int y) {
        int xroot = find(x);
        int yroot = find(y);
        if (Rank[xroot] < Rank[yroot])
            par[xroot] = yroot;
        else if (Rank[xroot] > Rank[yroot])
            par[yroot] = xroot;
        else {
            par[yroot] = xroot;
            Rank[xroot]++;
        }
    }

    // Function to compute minimum spanning tree
    static void computeMST() {
        // Sort edges in increasing order of weight
        Collections.sort(edges, Comparator.comparingInt(a -> a.first));

        // Go through all the edges
        for (int i = 0; i < m; i++) {
            int u = find(edges.get(i).second >> 10);
            int v = find(edges.get(i).second & ((1 << 10) - 1));

            if (u == v)
                continue;

            // Build minimum spanning tree and store minimum cost
            mst.add(edges.get(i));
            ans += edges.get(i).first;
            Union(u, v);
        }
    }

    // Function to find the cost of the minimum spanning tree
    static void modifiedKruskal(Pair x) {
        initialize();
        ans = 0;
        int sz = zeros.size();

        // Keep the edges which only have zero weights and remove all the other edges
        for (int i = 0; i < sz; i++) {
            Pair front = zeros.poll();
            if (front.first == x.first && front.second == x.second)
                continue;

            // Make union between the vertices of edges which have weight zero and keep them in the queue
            Union(front.first, front.second);
            zeros.add(front);
        }

        // Find the cost of the minimum spanning tree
        for (int i = 0; i < mst.size(); i++) {
            int u = find(mst.get(i).second >> 10);
            int v = find(mst.get(i).second & ((1 << 10) - 1));

            if (u == v)
                continue;

            ans += mst.get(i).first;
            Union(u, v);
        }
    }

    // Function to handle different queries
    static void query(int type, int u, int v) {
        // Update edge weight to 0
        if (type == 2) {
            zeros.add(new Pair(u, v));
            modifiedKruskal(new Pair(-1, -1));
        }
        // Restore edge weight to original value
        else if (type == 3) {
            zeros.add(new Pair(u, v));
            modifiedKruskal(new Pair(u, v));
        } else
            System.out.println(ans);
    }

    // Driver code
    public static void main(String[] args) {
        // Number of nodes and edges
        n = 4;
        m = 4;
        initialize();

        // Initialize data structures
        edges = new ArrayList<>();
        mst = new ArrayList<>();
        zeros = new LinkedList<>();

        // Add edges
        Add_edge(1, 2, 1);
        Add_edge(2, 3, 1);
        Add_edge(3, 4, 1);
        Add_edge(4, 1, 1);

        // Build the minimum spanning tree
        computeMST();

        // Execute queries
        query(2, 1, 2);
        query(1, 0, 0);
        query(3, 1, 2);
        query(1, 0, 0);
    }
}

Python3

# Python3 implementation of the approach
from collections import deque

N = 2005

# To store vertices, edges
# and the required answer
n, m, ans = 0, 0, 0

# To store parent and rank
par = [0] * N
Rank = [0] * N

# To store edges and the edges in MST
edges, mst = [], []

# To store the edges with weight zero
zeroes = deque()

# Function for initialize
def initialize():
    for i in range(n + 1):
        par[i] = i
        Rank[i] = 0

# Function to add edges
def add_edge(u: int, v: int, weight: int):
    edges.append((weight, (u, v)))

# Utility function to find set of an element i
# (uses path compression technique)
def find(x: int) -> int:
    if par[x] != x:
        par[x] = find(par[x])
    return par[x]

# Function that performs union of two sets x and y
# (uses union by rank)
def union(x: int, y: int):
    xroot = find(x)
    yroot = find(y)

    if Rank[xroot] < Rank[yroot]:
        par[xroot] = yroot
    elif Rank[xroot] > Rank[yroot]:
        par[yroot] = xroot
    else:
        par[yroot] = xroot
        Rank[xroot] += 1

# Function to compute minimum spanning tree
def compute_MST():
    global ans

    # Sort edges in increasing order of weight
    edges.sort()

    # Go through all the edges
    for i in range(m):
        u = find(edges[i][1][0])
        v = find(edges[i][1][1])

        if u == v:
            continue

        # Build minimum spanning tree
        # and store minimum cost
        mst.append(edges[i])
        ans += edges[i][0]
        union(u, v)

# Function to find the cost of minimum
# spanning tree
def modified_kruskal(x):
    global ans
    initialize()

    # Make answer zero
    ans = 0
    sz = len(zeroes)

    # Keep the edges which only have zero weights
    # and remove all the other edges
    for i in range(sz):
        front = zeroes[0]
        zeroes.popleft()

        if front[0] == x[0] and front[1] == x[1]:
            continue

        # Make union between the vertices of
        # edges which have weight zero and keep
        # them in queue
        union(front[0], front[1])
        zeroes.append(front)

    # Find the cost of the minimum spanning tree
    for i in range(len(mst)):
        u = find(mst[i][1][0])
        v = find(mst[i][1][1])

        if u == v:
            continue
        ans += mst[i][0]
        union(u, v)

# Function to handle different queries
def query(type: int, u=0, v=0):
    global ans

    # Update edge weight to 0
    if type == 2:

        # push edge in zeros
        zeroes.append((u, v))
        modified_kruskal((-1, -1))

    # Restore edge weight to original value
    elif type == 3:

        # push edge in zeros
        zeroes.append((u, v))
        modified_kruskal((u, v))
    else:
        print(ans)

# Driver Code
if __name__ == "__main__":

    # Number of nodes and edges
    n = 4
    m = 4
    initialize()

    # Add edges
    add_edge(1, 2, 1)
    add_edge(2, 3, 1)
    add_edge(3, 4, 1)
    add_edge(4, 1, 1)

    # Build the minimum spanning tree
    compute_MST()

    # Execute queries
    query(2, 1, 2)
    query(1)
    query(3, 1, 2)
    query(1)

# This code is contributed by
# sanjeev2552

using System;
using System.Collections.Generic;

class Program
{
    const int N = 2005;
    static int n, m, ans;
    static int[] par = new int[N];
    static int[] Rank = new int[N];
    static List<Tuple<int, Tuple<int, int>>> edges = new List<Tuple<int, Tuple<int, int>>>();
    static List<Tuple<int, Tuple<int, int>>> mst = new List<Tuple<int, Tuple<int, int>>>();
    static Queue<Tuple<int, int>> zeroes = new Queue<Tuple<int, int>>();

    static void Initialize()
    {
        // Initialize parent and rank arrays
        for (int i = 0; i <= n; i++)
        {
            par[i] = i;
            Rank[i] = 0;
        }
    }

    static void AddEdge(int u, int v, int weight)
    {
        // Add edge to the list
        edges.Add(new Tuple<int, Tuple<int, int>>(weight, new Tuple<int, int>(u, v)));
    }

    static int Find(int x)
    {
        // Find the root of the set using path compression
        if (par[x] != x)
        {
            par[x] = Find(par[x]);
        }
        return par[x];
    }

    static void Union(int x, int y)
    {
        // Union by rank
        int xroot = Find(x);
        int yroot = Find(y);

        if (Rank[xroot] < Rank[yroot])
        {
            par[xroot] = yroot;
        }
        else if (Rank[xroot] > Rank[yroot])
        {
            par[yroot] = xroot;
        }
        else
        {
            par[yroot] = xroot;
            Rank[xroot]++;
        }
    }

    static void ComputeMST()
    {
        ans = 0;
        // Sort edges in increasing order of weight
        edges.Sort();

        for (int i = 0; i < m; i++)
        {
            int u = Find(edges[i].Item2.Item1);
            int v = Find(edges[i].Item2.Item2);

            if (u == v)
            {
                // Skip if adding the edge creates a cycle
                continue;
            }

            // Build minimum spanning tree and update cost
            mst.Add(edges[i]);
            ans += edges[i].Item1;
            Union(u, v);
        }
    }

    static void ModifiedKruskal(Tuple<int, int> x)
    {
        Initialize();
        ans = 0;
        int sz = zeroes.Count;

        for (int i = 0; i < sz; i++)
        {
            Tuple<int, int> front = zeroes.Dequeue();

            if (front.Item1 == x.Item1 && front.Item2 == x.Item2)
            {
                continue;
            }

            Union(front.Item1, front.Item2);
            zeroes.Enqueue(front);
        }

        for (int i = 0; i < mst.Count; i++)
        {
            int u = Find(mst[i].Item2.Item1);
            int v = Find(mst[i].Item2.Item2);

            if (u == v)
            {
                continue;
            }
            ans += mst[i].Item1;
            Union(u, v);
        }
    }

    static void Query(int type, int u = 0, int v = 0)
    {
        if (type == 2)
        {
            // Update edge weight to 0
            zeroes.Enqueue(new Tuple<int, int>(u, v));
            ModifiedKruskal(new Tuple<int, int>(-1, -1));
        }
        else if (type == 3)
        {
            // Restore edge weight to original value
            zeroes.Enqueue(new Tuple<int, int>(u, v));
            ModifiedKruskal(new Tuple<int, int>(u, v));
        }
        else
        {
            // Print the cost of the minimum spanning tree
            Console.WriteLine(ans);
        }
    }

    static void Main()
    {
        n = 4;
        m = 4;
        Initialize();

        // Add edges
        AddEdge(1, 2, 1);
        AddEdge(2, 3, 1);
        AddEdge(3, 4, 1);
        AddEdge(4, 1, 1);

        // Build the minimum spanning tree
        ComputeMST();

        // Execute queries
        Query(2, 1, 2);
        Query(1);
        Query(3, 1, 2);
        Query(1);
    }
}

JavaScript

// JS implementation of the approach
const N = 2005;
let n = 0,
    m = 0,
    ans = 0;
    
// To store vertices, edges
// and the required answer
const par = new Array(N),
    Rank = new Array(N),
    edges = [],
    mst = [],
    zeroes = [];
    
// Function for initialize
function initialize() {
    for (let i = 0; i <= n; i++) {
        par[i] = i;
        Rank[i] = 0;
    }
}

// Function to add edges
function add_edge(u, v, weight) {
    edges.push([weight, [u, v]]);
}

// Utility function to find set of an element i
// (uses path compression technique)
function find(x) {
    if (par[x] !== x) {
        par[x] = find(par[x]);
    }
    return par[x];
}

// Function that performs union of two sets x and y
// (uses union by rank)
function union(x, y) {
    const xroot = find(x);
    const yroot = find(y);

    if (Rank[xroot] < Rank[yroot]) {
        par[xroot] = yroot;
    } else if (Rank[xroot] > Rank[yroot]) {
        par[yroot] = xroot;
    } else {
        par[yroot] = xroot;
        Rank[xroot]++;
    }
}

// Function to compute minimum spanning tree
function compute_MST() {

     // Sort edges in increasing order of weight
    edges.sort((a, b) => a[0] - b[0]);
    
       // Go through all the edges
    for (let i = 0; i < m; i++) {
        const u = find(edges[i][1][0]);
        const v = find(edges[i][1][1]);
        if (u === v) continue;
        
         // Build minimum spanning tree
       // and store minimum cost
        mst.push(edges[i]);
        ans += edges[i][0];
        union(u, v);
    }
}

// Function to find the cost of minimum
// spanning tree
function modified_kruskal(x) {
    initialize();
    
     // Make answer zero
    ans = 0;
    const sz = zeroes.length;
    
       // Keep the edges which only have zero weights
   // and remove all the other edges
    for (let i = 0; i < sz; i++) {
        const front = zeroes.shift();
        if (front[0] === x[0] && front[1] === x[1]) continue;
        
          // Make union between the vertices of
       // edges which have weight zero and keep
       // them in queue
        union(front[0], front[1]);
        zeroes.push(front);
    }
    
      // Find the cost of the minimum spanning tree
    for (let i = 0; i < mst.length; i++) {
        const u = find(mst[i][1][0]);
        const v = find(mst[i][1][1]);
        if (u === v) continue;
        ans += mst[i][0];
        union(u, v);
    }
}

// Function to handle different queries
function query(type, u = 0, v = 0) {

       // Update edge weight to 0
    if (type === 2) {
    
           // push edge in zeros
        zeroes.push([u, v]);
        modified_kruskal([-1, -1]);
        
    }// Restore edge weight to original value
    else if (type === 3) {
           // push edge in zeros
        zeroes.push([u, v]);
        modified_kruskal([u, v]);
    } else {
        console.log(ans);
    }
}

 // Number of nodes and edges
n = 4;
m = 4;
initialize();
add_edge(1, 2, 1);
add_edge(2, 3, 1);
add_edge(3, 4, 1);
add_edge(4, 1, 1);

 // Build the minimum spanning tree
compute_MST();

query(2, 1, 2);
query(1);
query(3, 1, 2);
query(1);

// This code is contributed by lokeshpotta20.

Output

2
3

Time Complexity: O(N) per query, where N is the total number of nodes in the graph.
Auxiliary Space: O(N)

Find the weight of the minimum spanning tree

Explore