From 8467b527ba25fff9882d5c1f0067595b1f70466a Mon Sep 17 00:00:00 2001
From: Timofei Rusanov <92477489+taaae@users.noreply.github.com>
Date: Thu, 26 Dec 2024 11:32:16 +0100
Subject: [PATCH] Fixed dijkstra implementation bug: improved time from
 exponential to O(E * logV) (#848)

Fixed dijkstra implementation: the previous implementation incorrectly prioritized unvisited nodes with the smallest vertex instead of the smallest weight.
---
 src/graph/dijkstra.rs | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/graph/dijkstra.rs b/src/graph/dijkstra.rs
index b27dfa8c34f..8cef293abe7 100644
--- a/src/graph/dijkstra.rs
+++ b/src/graph/dijkstra.rs
@@ -1,10 +1,10 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::ops::Add;
 
 type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;
 
 // performs Dijsktra's algorithm on the given graph from the given start
-// the graph is a positively-weighted undirected graph
+// the graph is a positively-weighted directed graph
 //
 // returns a map that for each reachable vertex associates the distance and the predecessor
 // since the start has no predecessor but is reachable, map[start] will be None
@@ -17,17 +17,17 @@ pub fn dijkstra<V: Ord + Copy, E: Ord + Copy + Add<Output = E>>(
     start: V,
 ) -> BTreeMap<V, Option<(V, E)>> {
     let mut ans = BTreeMap::new();
-    let mut prio = BTreeMap::new();
+    let mut prio = BTreeSet::new();
 
     // start is the special case that doesn't have a predecessor
     ans.insert(start, None);
 
     for (new, weight) in &graph[&start] {
         ans.insert(*new, Some((start, *weight)));
-        prio.insert(*new, *weight);
+        prio.insert((*weight, *new));
     }
 
-    while let Some((vertex, path_weight)) = prio.pop_first() {
+    while let Some((path_weight, vertex)) = prio.pop_first() {
         for (next, weight) in &graph[&vertex] {
             let new_weight = path_weight + *weight;
             match ans.get(next) {
@@ -37,8 +37,12 @@ pub fn dijkstra<V: Ord + Copy, E: Ord + Copy + Add<Output = E>>(
                 Some(None) => {}
                 // the new path is shorter, either new was not in ans or it was farther
                 _ => {
-                    ans.insert(*next, Some((vertex, new_weight)));
-                    prio.insert(*next, new_weight);
+                    if let Some(Some((_, prev_weight))) =
+                        ans.insert(*next, Some((vertex, new_weight)))
+                    {
+                        prio.remove(&(prev_weight, *next));
+                    }
+                    prio.insert((new_weight, *next));
                 }
             }
         }