Shortest Distance After Road Addition Queries II solution and explanation

from collections import deque

def shortest_path(n, queries):
    shortest_path_lengths = []
    adjacency_list = {i: set() for i in range(n)}
    for i in range(n - 1):
        adjacency_list[i].add(i + 1)
    
    def bfs_dist():
        queue = deque([0])
        distances = {i: float('inf') for i in range(n)}
        distances[0] = 0
        visited = set([0])
        while queue:
            city = queue.popleft()
            current_distance = distances[city]
            for neighbor in adjacency_list[city]:
                if neighbor not in visited:
                    queue.append(neighbor)
                    visited.add(neighbor)
                    distances[neighbor] = current_distance + 1
                    if neighbor == n - 1:
                        return distances[n - 1]
        return distances[n - 1]

    for u, v in queries:
        adjacency_list[u].add(v)
        shortest_path_lengths.append(bfs_dist())
    
    return shortest_path_lengths

n = 5
queries = [[2, 4], [0, 2], [0, 4]]
print(shortest_path(n, queries))

import java.util.*;

public class ShortestPath {
    static List<Integer> shortestPath(int n, int[][] queries) {
        List<Integer> result = new ArrayList<>();
        Map<Integer, List<Integer>> graph = new HashMap<>();
        for (int i = 0; i < n; i++) {
            graph.put(i, new ArrayList<>());
        }
        for (int i = 0; i < n - 1; i++) {
            graph.get(i).add(i + 1);
        }
        for (int[] q : queries) {
            graph.get(q[0]).add(q[1]);
            result.add(bfsDist(graph, n));
        }
        return result;
    }

    private static int bfsDist(Map<Integer, List<Integer>> graph, int n) {
        Queue<Integer> queue = new LinkedList<>();
        boolean[] visited = new boolean[n];
        int[] distances = new int[n];
        Arrays.fill(distances, Integer.MAX_VALUE);
        queue.add(0);
        visited[0] = true;
        distances[0] = 0;

        while (!queue.isEmpty()) {
            int city = queue.poll();
            for (int neighbor : graph.get(city)) {
                if (!visited[neighbor]) {
                    visited[neighbor] = true;
                    distances[neighbor] = distances[city] + 1;
                    queue.add(neighbor);
                    if (neighbor == n - 1) {
                        return distances[n - 1];
                    }
                }
            }
        }

        return distances[n - 1];
    }

    public static void main(String[] args) {
        int n = 5;
        int[][] queries = {{2, 4}, {0, 2}, {0, 4}};
        System.out.println(shortestPath(n, queries));
    }
}

#include <iostream>
#include <vector>
#include <numeric>

using namespace std;

class DSU {
    vector<int> parent, rank;

public:
    DSU(int n) : parent(n), rank(n, 0) {
        iota(parent.begin(), parent.end(), 0);
    }

    int find(int u) {
        if (u != parent[u]) {
            parent[u] = find(parent[u]);
        }
        return parent[u];
    }

    bool unionSets(int u, int v) {
        int pu = find(u);
        int pv = find(v);
        if (pu == pv) return false;
        if (rank[pu] < rank[pv]) {
            parent[pu] = pv;
        } else if (rank[pu] > rank[pv]) {
            parent[pv] = pu;
        } else {
            parent[pv] = pu;
            rank[pu]++;
        }
        return true;
    }
};

vector<int> shortestPathC(int n, vector<vector<int>>& queries) {
    DSU dsu(n);
    vector<int> answer;
    for (auto& query : queries) {
        dsu.unionSets(query[0], query[1]);
        if (dsu.find(0) == dsu.find(n - 1)) {
            answer.push_back(1);
        } else {
            int additional_path_length = 0;
            for (int i = 0; i < n - 1; ++i) {
                if (dsu.find(i) != dsu.find(i + 1)) {
                    additional_path_length++;
                }
            }
            answer.push_back(additional_path_length + 1);
        }
    }
    return answer;
}

int main() {
    int n = 5;
    vector<vector<int>> queries = {{2, 4}, {0, 2}, {0, 4}};
    vector<int> result = shortestPathC(n, queries);
    for (int len : result) {
        cout << len << ' ';
    }
    return 0;
}

class DSU {
    constructor(n) {
        this.parent = Array.from(Array(n).keys());
        this.rank = Array(n).fill(0);
    }

    find(u) {
        if (this.parent[u] !== u) this.parent[u] = this.find(this.parent[u]);
        return this.parent[u];
    }

    union(u, v) {
        let pu = this.find(u);
        let pv = this.find(v);
        if (pu === pv) return false;
        if (this.rank[pu] < this.rank[pv]) {
            this.parent[pu] = pv;
        } else if (this.rank[pu] > this.rank[pv]) {
            this.parent[pv] = pu;
        } else {
            this.parent[pv] = pu;
            this.rank[pu]++;
        }
        return true;
    }
}

function shortestPathJS(n, queries) {
    const dsu = new DSU(n);
    const result = [];
    for (const [u, v] of queries) {
        dsu.union(u, v);
        if (dsu.find(0) === dsu.find(n - 1)) {
            result.push(1);
        } else {
            let extraEdges = 0;
            for (let i = 0; i < n - 1; i++) {
                if (dsu.find(i) !== dsu.find(i + 1)) {
                    extraEdges++;
                }
            }
            result.push(extraEdges + 1);
        }
    }
    return result;
}

console.log(shortestPathJS(5, [[2, 4], [0, 2], [0, 4]]));

function shortestDistanceAfterQueries(n: number, queries: number[][]): number[] {
    const nxt: number[] = Array.from({ length: n - 1 }, (_, i) => i + 1);
    const ans: number[] = [];
    let cnt = n - 1;
    for (const [u, v] of queries) {
        if (nxt[u] && nxt[u] < v) {
            let i = nxt[u];
            while (i < v) {
                --cnt;
                [nxt[i], i] = [0, nxt[i]];
            }
            nxt[u] = v;
        }
        ans.push(cnt);
    }
    return ans;
}

func shortestDistanceAfterQueries(n int, queries [][]int) (ans []int) {
	nxt := make([]int, n-1)
	for i := range nxt {
		nxt[i] = i + 1
	}
	cnt := n - 1
	for _, q := range queries {
		u, v := q[0], q[1]
		if nxt[u] > 0 && nxt[u] < v {
			i := nxt[u]
			for i < v {
				cnt--
				nxt[i], i = 0, nxt[i]
			}
			nxt[u] = v
		}
		ans = append(ans, cnt)
	}
	return
}