CLRS-Lean
Chapters 1–35 · CLRS Fourth Edition
Chapter 23. All-Pairs Shortest Paths
23.3. Johnson's Algorithm for Sparse Graphs

Imports

import Mathlib
import CLRSLean.FourthEdition.Chapter_22.Section_22_1_Bellman_Ford
import CLRSLean.FourthEdition.Chapter_22.Section_22_3_Dijkstra

23.3. Johnson's Algorithm for Sparse Graphs

Johnson's algorithm computes all-pairs shortest paths in a weighted directed graph with no negative-weight cycles. It works in three stages:

Potential via Bellman-Ford: add a new source vertex s with zero-weight edges to every vertex, run Bellman-Ford from s, and set h(v) = delta(s, v). This construction assumes global NoNegCycle; it does not return a negative-cycle failure or implement an abort branch.
Reweighting: define a new weight function w^(u,v) = w(u,v) + h(u) - h(v). By the triangle inequality of shortest paths, w^ >= 0 for every edge, so Dijkstra's algorithm can be used from every vertex.
|V| x Dijkstra: run Dijkstra from every vertex on the reweighted graph and recover the true distances via delta(u,v) = delta^(u,v) - h(u) + h(v).

Main results:

Theorem reweightedWeight_nonneg: reweighted edge weights are nonnegative.
Theorem reweightedWalkWeight_eq: w^(p) = w(p) + h(u) - h(v) for any walk p from u to v (telescoping property).
Theorem reweighted_isShortestDist: reweighted shortest distances equal original distances shifted by h(u) - h(v).
Lemma noNegCycle_johnsonAugmentedGraph: the augmented graph has no negative cycles iff the original graph has none.
Lemma johnsonPotential_triangle: the Bellman-Ford potential satisfies h(v) ≤ h(u) + w(u, v) for every edge (u, v).
Theorem johnsonDist_isShortestDist: CLRS Theorem 23.5 — end-to-end correctness of Johnson's algorithm: johnsonDist computes the exact all-pairs shortest-path distances.
Lemma johnsonAugmentedGraph_edges_card: the augmented graph has |V| + |E| edges.
Theorem johnsonCost_eq: an independent heap-backend budget |V|·(|V| + |E|)·(log₂|V| + 2). This formula alone does not count a concrete queue or stored distance table.
Theorem johnsonCost_le: the 2·|V|·(|V| + |E|)·(log₂|V| + 1) O-bound corollary.

The section is complete: the augmented-graph potential construction, triangle inequality, reweighting nonnegativity, end-to-end Johnson correctness, and the running-time bound are all proved.

              namespace CLRSnamespace Chapter24open Finsetnamespace WeightedGraphvariable {V : Type*} [Fintype V] [DecidableEq V] (G : WeightedGraph V)

Johnson's algorithm preliminary definitions

The augmented graph for Johnson's algorithm: add a fresh source vertex none with zero-weight edges to every original vertex.

All original edges are preserved, with the original weights. There are no edges entering none, so any negative cycle in the original graph remains a negative cycle (and no new negative cycles are introduced).

noncomputable def johnsonAugmentedGraph : WeightedGraph (Option V) :=
  { edges := (Finset.image (fun (v : V) => (none, some v)) Finset.univ) ∪
             (Finset.image (fun ((u, v) : V × V) => (some u, some v)) G.edges)
  , w := fun u v =>
      match u, v with
      | none, some Variable name `v` is not explicitly referenced.

The binding can be removed (if unused) or named `_` (if used implicitly).

Note: This linter can be disabled with `set_option linter.unusedVariables false`v => 0
      | some u, some v => G.w u v
      | _, _ => 0
  }

@[simp] theorem mem_edges_johnsonAugmentedGraph_source (v : V) :
    (none, some v) ∈ G.johnsonAugmentedGraph.edges := by
  unfold johnsonAugmentedGraph; simp

@[simp] theorem mem_edges_johnsonAugmentedGraph_edge (u v : V) :
    (some u, some v) ∈ G.johnsonAugmentedGraph.edges ↔ (u, v) ∈ G.edges := by
  unfold johnsonAugmentedGraph; simp

There are no edges entering none in the augmented graph.

theorem no_incoming_to_none_johnsonAugmentedGraph (u : Option V) :
    (u, none) ∉ G.johnsonAugmentedGraph.edges := by
  unfold johnsonAugmentedGraph; simp

Reweighting with a potential function

The reweighted weight function w^(u,v) = w(u,v) + h(u) - h(v), where h : V -> RR is a potential function (typically delta(s, .) from Bellman-Ford).

def reweightedWeight (h : V → ℝ) (u v : V) : ℝ :=
  G.w u v + h u - h v

@[simp] theorem reweightedWeight_eq (h : V → ℝ) (u v : V) :
    G.reweightedWeight h u v = G.w u v + h u - h v := rfl

The reweighted graph G^ has the same edge set as G but with the reweighted weight function w^.

noncomputable def reweightedGraph (h : V → ℝ) : WeightedGraph V :=
  { edges := G.edges
  , w := G.reweightedWeight h
  }

@[simp] theorem edges_reweightedGraph (h : V → ℝ) :
    (G.reweightedGraph h).edges = G.edges := rfl

@[simp] theorem w_reweightedGraph (h : V → ℝ) (u v : V) :
    (G.reweightedGraph h).w u v = G.reweightedWeight h u v := rfl

Telescoping property. For any walk p from u to v, the reweighted walk weight equals the original walk weight plus h(u) - h(v).


theorem reweightedWalkWeight_eq (h : V → ℝ) (u v : V) (p : List V)
    (hp : G.IsWalkFrom u v p) : walkWeight (G.reweightedWeight h) p = walkWeight G.w p + h u - h v := by
  induction p generalizing u v with
  | nil => exact absurd rfl hp.ne_nil
  | cons a as ih =>
    have ha_u : a = u := by
      have hh := hp.head; simpa using hh
    rw [ha_u]
    rw [ha_u] at hp
    cases as with
    | nil =>
      have hv_u : v = u := by
        have hl := hp.last; simp at hl; exact hl.symm
      rw [hv_u]; simp [walkWeight, This simp argument is unused:
  reweightedWeight

Hint: Omit it from the simp argument list.
  simp [walkWeight,̵ ̵r̵e̵w̵e̵i̵g̵h̵t̵e̵d̵W̵e̵i̵g̵h̵t̵]

Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`reweightedWeight]
    | cons b bs =>
      have h_chain : List.IsChain G.Adj (u :: b :: bs) := hp.chain
      have h_chain_rest : List.IsChain G.Adj (b :: bs) := by
        cases h_chain with
        | cons_cons _ htail => exact htail
      have h_walk_rest : G.IsWalkFrom b v (b :: bs) :=
        ⟨h_chain_rest, by simp, by simpa using hp.last⟩
      simp [walkWeight, reweightedWeight]
      rw [ih b v h_walk_rest]
      ring

Nonnegativity of reweighted weights. If the potential h satisfies the triangle inequality h(v) <= h(u) + w(u, v) for every edge (u, v), then every reweighted edge weight is nonnegative.

theorem reweightedWeight_nonneg (h : V → ℝ)
    (h_triangle : ∀ u v, (u, v) ∈ G.edges → h v ≤ h u + G.w u v) :
    ∀ u v, (u, v) ∈ G.edges → 0 ≤ G.reweightedWeight h u v := by
  intro u v h_edge
  dsimp [reweightedWeight]
  have hineq : h v ≤ h u + G.w u v := h_triangle u v h_edge
  linarith

Walk equivalence: original graph ↔ reweighted graph



lemma IsWalkFrom_reweighted_iff (h : V → ℝ) (u v : V) (p : List V) :
    (G.reweightedGraph h).IsWalkFrom u v p ↔ G.IsWalkFrom u v p := by
  have h_adj_eq : (G.reweightedGraph h).Adj = G.Adj := by
    ext x y; simp [WeightedGraph.Adj, edges_reweightedGraph]
  constructor
  · intro hw; rcases hw with ⟨hc, hh, hl⟩; refine ⟨?_, hh, hl⟩; rw [← h_adj_eq]; exact hc
  · intro hw; rcases hw with ⟨hc, hh, hl⟩; refine ⟨?_, hh, hl⟩; rw [h_adj_eq]; exact hc

Lift telescoping property to WithTop ℝ

lemma reweightedWalkWeight_eq_withtop (h : V → ℝ) (u v : V) (p : List V)
    (hp : G.IsWalkFrom u v p) :
    (walkWeight (G.reweightedWeight h) p : WithTop ℝ) =
      (walkWeight G.w p : WithTop ℝ) + (h u : WithTop ℝ) - (h v : WithTop ℝ) := by
  have h_rw := reweightedWalkWeight_eq G h u v p hp
  calc
    (walkWeight (G.reweightedWeight h) p : WithTop ℝ) =
        ((walkWeight G.w p + h u - h v : ℝ) : WithTop ℝ) := by exact_mod_cast h_rw
    _ = (walkWeight G.w p : WithTop ℝ) + (h u : WithTop ℝ) - (h v : WithTop ℝ) := by simp

WithTop helpers

private lemma add_sub_cancel {a b c : WithTop ℝ} (hc : c ≠ ⊤) (h : a + c ≤ b + c) : a ≤ b :=
  (WithTop.add_le_add_iff_right hc).mp h



private lemma add_sub_assoc {a b c : WithTop ℝ} (hb : b ≠ ⊤) (hc : c ≠ ⊤) :
    (a + b) - c = a + (b - c) := by
  have ⟨b', hb'⟩ := Option.ne_none_iff_exists'.mp hb
  have ⟨c', hc'⟩ := Option.ne_none_iff_exists'.mp hc
  subst hb' hc'
  by_cases ha : a = ⊤
  · subst ha
    calc
      ((⊤ : WithTop ℝ) + (b' : WithTop ℝ)) - (c' : WithTop ℝ) = (⊤ : WithTop ℝ) - (c' : WithTop ℝ) := by simp
      _ = (⊤ : WithTop ℝ) := by simp
      _ = (⊤ : WithTop ℝ) + ((b' : WithTop ℝ) - (c' : WithTop ℝ)) := by rw [WithTop.top_add]
  · have ⟨a', ha'⟩ := Option.ne_none_iff_exists'.mp ha
    subst ha'
    -- ℝ equality: (a'+b'-c') = (a'+(b'-c')) in ℝ, lifted to WithTop ℝ
    calc
      ((a' : ℝ) : WithTop ℝ) + ((b' : ℝ) : WithTop ℝ) - ((c' : ℝ) : WithTop ℝ) =
          (((a' + b' - c' : ℝ) : ℝ) : WithTop ℝ) := by simp
      _ = (((a' + (b' - c') : ℝ) : ℝ) : WithTop ℝ) := by rw [show (a' + b' - c' : ℝ) = (a' + (b' - c') : ℝ) by ring]
      _ = ((a' : ℝ) : WithTop ℝ) + (((b' : ℝ) : WithTop ℝ) - ((c' : ℝ) : WithTop ℝ)) := by simp

Shortest-path preservation under reweighting

Shortest-path preservation. For any potential h, the reweighted shortest distance equals the original shortest distance shifted by h(u) - h(v). This holds without any feasibility assumption on h.


theorem reweighted_isShortestDist (h : V → ℝ) (u v : V) (d : WithTop ℝ) :
    (G.reweightedGraph h).IsShortestDist u v
      (d + (h u : WithTop ℝ) - (h v : WithTop ℝ)) ↔
    G.IsShortestDist u v d := by
  have h_fin_hu : (h u : WithTop ℝ) ≠ ⊤ := by simp
  have h_fin_hv : (h v : WithTop ℝ) ≠ ⊤ := by simp
  have h_fin_diff : ((h u : WithTop ℝ) - (h v : WithTop ℝ)) ≠ ⊤ := by simp
  have h_walk_iff := IsWalkFrom_reweighted_iff G h
  have h_rw_eq := reweightedWalkWeight_eq_withtop G h
  constructor
  · intro h_rsd; rcases h_rsd with ⟨h_lower, h_att⟩; constructor
    · intro p hp
      have hp_hat : (G.reweightedGraph h).IsWalkFrom u v p := (h_walk_iff u v p).mpr hp
      have h_bound := h_lower p hp_hat
      have h_rw' := h_rw_eq u v p hp
      have h_bound' : d + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) ≤
          (walkWeight G.w p : WithTop ℝ) + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) := by
        -- Regroup using add_sub_assoc and use h_bound + h_rw'
        simpa [add_sub_assoc h_fin_hu h_fin_hv] using
          calc
            d + (h u : WithTop ℝ) - (h v : WithTop ℝ) ≤
                (walkWeight (G.reweightedWeight h) p : WithTop ℝ) := h_bound
            _ = (walkWeight G.w p : WithTop ℝ) + (h u : WithTop ℝ) - (h v : WithTop ℝ) := h_rw'
      exact add_sub_cancel h_fin_diff h_bound'
    · rcases h_att with (h_dtop | ⟨p, hp_hat, hpw⟩)
      · -- h_dtop: (d + h_u - h_v) = ⊤ in reweighted graph
        -- Need: d = ⊤ in original graph.  Since h_u, h_v are finite,
        -- (d + h_u) - h_v = ⊤ implies d + h_u = ⊤, which implies d = ⊤.
        by_cases hd_top : d = ⊤
        · exact Or.inl hd_top
        · exfalso
          -- d ≠ ⊤, h_u ≠ ⊤, h_v ≠ ⊤ → d + h_u - h_v ≠ ⊤, contradicting h_dtop
          have h_sum_fin : d + (h u : WithTop ℝ) - (h v : WithTop ℝ) ≠ ⊤ := by
            -- All three are finite ℝ values, so the sum is finite
            -- Proof: case analysis to extract the ℝ values
            rcases Option.ne_none_iff_exists'.mp hd_top with ⟨d', hd'⟩
            rcases Option.ne_none_iff_exists'.mp h_fin_hu with ⟨hu', hhu'⟩
            rcases Option.ne_none_iff_exists'.mp h_fin_hv with ⟨hv', hhv'⟩
            rw [hd', hhu', hhv']
            simp
          exact h_sum_fin h_dtop
      · right
        have hp : G.IsWalkFrom u v p := (h_walk_iff u v p).mp hp_hat
        have h_rw' := h_rw_eq u v p hp
        have hpw' : (walkWeight G.w p : WithTop ℝ) = d := by
          have h_eq : (walkWeight G.w p : WithTop ℝ) + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) =
              d + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) := by
            calc
              (walkWeight G.w p : WithTop ℝ) + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) =
                  (walkWeight G.w p : WithTop ℝ) + (h u : WithTop ℝ) - (h v : WithTop ℝ) := by
                rw [add_sub_assoc h_fin_hu h_fin_hv]
              _ = (walkWeight (G.reweightedWeight h) p : WithTop ℝ) := by symm; exact h_rw'
              _ = d + (h u : WithTop ℝ) - (h v : WithTop ℝ) := hpw
              _ = d + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) := by
                rw [add_sub_assoc h_fin_hu h_fin_hv]
          apply le_antisymm
          · apply add_sub_cancel h_fin_diff; exact h_eq.le
          · apply add_sub_cancel h_fin_diff; exact h_eq.ge
        refine ⟨p, hp, hpw'⟩
  · intro h_sd; rcases h_sd with ⟨h_lower, h_att⟩; constructor
    · intro p hp_hat
      have hp : G.IsWalkFrom u v p := (h_walk_iff u v p).mp hp_hat
      have h_rw' := h_rw_eq u v p hp
      rw [add_sub_assoc h_fin_hu h_fin_hv]
      have h_add_both : d + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) ≤
          (walkWeight G.w p : WithTop ℝ) + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) := by
        gcongr; exact h_lower p hp
      calc
        d + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) ≤
            (walkWeight G.w p : WithTop ℝ) + ((h u : WithTop ℝ) - (h v : WithTop ℝ)) := h_add_both
        _ = (walkWeight G.w p : WithTop ℝ) + (h u : WithTop ℝ) - (h v : WithTop ℝ) := by
          rw [add_sub_assoc h_fin_hu h_fin_hv]
        _ = (walkWeight (G.reweightedWeight h) p : WithTop ℝ) := h_rw'.symm
    · rcases h_att with (h_dtop | ⟨p, hp, hpw⟩)
      · -- h_dtop: d = ⊤ in original graph
        -- Need: d + h_u - h_v = ⊤ in reweighted graph
        refine Or.inl ?_
        rw [h_dtop]; simp
      · right
        have hp_hat : (G.reweightedGraph h).IsWalkFrom u v p := (h_walk_iff u v p).mpr hp
        have h_rw' := h_rw_eq u v p hp
        refine ⟨p, hp_hat, ?_⟩
        calc
          (walkWeight (G.reweightedWeight h) p : WithTop ℝ) =
              (walkWeight G.w p : WithTop ℝ) + (h u : WithTop ℝ) - (h v : WithTop ℝ) := h_rw'
          _ = d + (h u : WithTop ℝ) - (h v : WithTop ℝ) := by rw [hpw]

Negative-cycle equivalence for the augmented graph

No edge in johnsonAugmentedGraph targets none.

lemma adj_target_ne_none (u v : Option V)
    (h_adj : G.johnsonAugmentedGraph.Adj u v) : v ≠ none := by
  intro h_eq; subst h_eq
  have h_edge : (u, none) ∈ G.johnsonAugmentedGraph.edges := h_adj
  exact G.no_incoming_to_none_johnsonAugmentedGraph u h_edge

If a chain in johnsonAugmentedGraph does not start at none, then none never appears in the chain.


lemma chain_no_none (l : List (Option V))
    (h_chain : List.IsChain G.johnsonAugmentedGraph.Adj l)
    (h_head_ne_none : l.head? ≠ some none) : none ∉ l := by
  induction l with
  | nil => simp
  | cons x l' ih =>
    rw [List.isChain_cons] at h_chain
    rcases h_chain with ⟨h_adj, h_chain_tail⟩
    have hx_ne_none : x ≠ none := by
      intro h_eq; subst h_eq; apply h_head_ne_none; simp
    have h_l'_head_ne_none : l'.head? ≠ some none := by
      intro h_eq
      have h_mem : none ∈ l'.head? := by rw [h_eq]; simp
      have h_adj_x_none : G.johnsonAugmentedGraph.Adj x none := h_adj none h_mem
      exact adj_target_ne_none G x none h_adj_x_none rfl
    have h_none_notin_l' : none ∉ l' := ih h_chain_tail h_l'_head_ne_none
    intro h
    cases h with
    | head _ => exact hx_ne_none rfl
    | tail _ h_mem => exact h_none_notin_l' h_mem

The only walk from none to none in the augmented graph is [none].


lemma walk_from_none_to_none_singleton (c : List (Option V))
    (hc : G.johnsonAugmentedGraph.IsWalkFrom none none c) : c = [none] := by
  have h_last : c.getLast? = some none := hc.last
  rcases List.getLast?_eq_some_iff.mp h_last with ⟨c₁, hc_eq⟩
  subst hc_eq
  by_cases h_c₁_empty : c₁ = []
  · subst h_c₁_empty; simp
  · have h_chain : List.IsChain G.johnsonAugmentedGraph.Adj (c₁ ++ [none]) := hc.chain
    rw [List.isChain_append] at h_chain
    rcases h_chain with ⟨_, _, h_adj_cond⟩
    have h_c₁_last_some : c₁.getLast? ≠ none := by
      intro h_eq; apply h_c₁_empty
      exact (List.getLast?_eq_none_iff.mp h_eq)
    rcases Option.ne_none_iff_exists'.mp h_c₁_last_some with ⟨last_c₁, h_last_c₁⟩
    have h_adj_last_none : G.johnsonAugmentedGraph.Adj last_c₁ none :=
      h_adj_cond last_c₁ (by rw [h_last_c₁]; simp) none (by simp)
    exfalso
    exact (adj_target_ne_none G last_c₁ none h_adj_last_none) rfl

Extend a walk by a single edge u → v at the end.


lemma IsWalkFrom.append_step (hp : G.IsWalkFrom s u p) (h_edge : G.Adj u v) :
    G.IsWalkFrom s v (p ++ [v]) := by
  have h_chain : List.IsChain G.Adj (p ++ [v]) :=
    List.IsChain.append hp.chain (List.isChain_singleton v)
      (by
        intro x hx_last
        -- hx_last : p.getLast? = some x (by definition of ∈ for Option)
        rw [hp.last] at hx_last
        -- hx_last : some u = some x
        have hx_eq : x = u := (Option.some_inj.mp hx_last).symm
        subst hx_eq
        intro y hy_head
        -- hy_head : [v].head? = some y (by definition of ∈ for Option)
        -- [v].head? = some v (by simp)
        have hy_eq : y = v := by
          -- hy_head: [v].head? = some y, and [v].head? = some v
          -- So some y = some v, hence y = v
          have hh : [v].head? = some y := hy_head
          have hv : [v].head? = some v := by simp
          rw [hv] at hh
          -- hh: some v = some y
          exact (Option.some_inj.mp hh).symm
        subst hy_eq
        exact h_edge)
  have h_head : (p ++ [v]).head? = some s := by
    have hp_ne_nil : p ≠ [] := hp.ne_nil
    have hp_head_val : p.head? = some s := hp.head
    cases p with
    | nil => exact absurd rfl hp_ne_nil
    | cons a as =>
      -- p = a :: as, so (a :: as ++ [v]).head? = some a
      -- and hp.head gives some a = some s
      have h_head_a : (a :: as).head? = some a := by simp
      rw [h_head_a] at hp_head_val
      -- hp_head_val : some a = some s
      simp [hp_head_val]
  have h_last : (p ++ [v]).getLast? = some v := by simp
  exact ⟨h_chain, h_head, h_last⟩

The weight of a walk extended by a single edge.

lemma walkWeight_append_step (hp : G.IsWalkFrom s u p) (Variable name `h_edge` is not explicitly referenced.

The binding can be removed (if unused) or named `_` (if used implicitly).

Note: This linter can be disabled with `set_option linter.unusedVariables false`h_edge : G.Adj u v) :
    (walkWeight G.w (p ++ [v]) : WithTop ℝ) =
      (walkWeight G.w p : WithTop ℝ) + (G.w u v : WithTop ℝ) := by
  rcases List.getLast?_eq_some_iff.mp hp.last with ⟨q, hq⟩
  subst hq
  simp [walkWeight_concat G.w q u v]

An auxiliary function that strips Option.some from an element.

private def unsome (default : V) (x : Option V) : V :=
  match x with | some v => v | none => default

For a none-free list, the chain in G' lifts to a chain in G.


private lemma chain_none_free_map (default : V) (l : List (Option V))
    (h_chain : List.IsChain G.johnsonAugmentedGraph.Adj l)
    (h_no_none : none ∉ l) :
    List.IsChain G.Adj (l.map (unsome default)) := by
  induction l with
  | nil => exact List.IsChain.nil
  | cons x l' ih =>
    rw [List.isChain_cons] at h_chain
    rcases h_chain with ⟨h_adj, h_chain_tail⟩
    have hx_ne_none : x ≠ none := by
      intro h_eq; subst h_eq; exact h_no_none (by simp)
    have hl'_no_none : none ∉ l' := by
      intro h; exact h_no_none (by simp [h])
    -- Determine x (must be some a since ≠ none)
    obtain ⟨a, hx_eq⟩ := Option.ne_none_iff_exists'.mp hx_ne_none
    subst hx_eq
    -- Now x = some a
    have ih_map : List.IsChain G.Adj (l'.map (unsome default)) :=
      ih h_chain_tail hl'_no_none
    rw [show (some a :: l').map (unsome default) = a :: l'.map (unsome default) by
      simp [unsome]]
    apply List.IsChain.cons ih_map
    intro y hy
    -- hy : y ∈ (l'.map (unsome default)).head?
    -- We need G.Adj a y. Analyze l'.
    cases l' with
    | nil => simp at hy
    | cons z l'' =>
      -- z must be some b (since none ∉ l')
      have hz_ne_none : z ≠ none := by
        intro h_eq; subst h_eq; exact hl'_no_none (by simp)
      obtain ⟨b, hz_eq⟩ := Option.ne_none_iff_exists'.mp hz_ne_none
      subst hz_eq
      -- Now l' = some b :: l''
      -- The head of (some b :: l'').map (unsome default) is b
      have h_head_map : ((some b :: l'').map (unsome default)).head? = some b := by
        simp [unsome]
      rw [h_head_map] at hy
      -- Now hy : y ∈ some b, i.e. some b = some y, so b = y
      have hy_eq : b = y := Option.some_inj.mp hy
      subst hy_eq
      -- Need G.Adj a b. From h_adj: G'.Adj (some a) (some b)
      have h_adj' : G.johnsonAugmentedGraph.Adj (some a) (some b) :=
        h_adj (some b) (by simp)
      exact (G.mem_edges_johnsonAugmentedGraph_edge a b).mp h_adj'

For a none-free list, the walk weights in G' and G agree.


private lemma walkWeight_none_free_eq (default : V) (c : List (Option V))
    (h_no_none : none ∉ c) :
    walkWeight G.johnsonAugmentedGraph.w c =
      walkWeight G.w (c.map (unsome default)) := by
  induction c with
  | nil => simp [walkWeight]
  | cons x cs ih =>
    cases x with
    | none => exfalso; exact h_no_none (by simp)
    | some a =>
      cases cs with
      | nil =>
        -- c = [some a]; both sides are 0
        simp [walkWeight]
      | cons y rest =>
        have h_cs_no_none : none ∉ (y :: rest) := by
          intro h; apply h_no_none; simp [h]
        cases y with
        | none => exfalso; exact h_cs_no_none (by simp)
        | some b =>
          -- c = some a :: some b :: rest
          -- walkWeight G'.w c = G.w a b + walkWeight G'.w (some b :: rest)
          -- walkWeight G.w (c.map f) = G.w a b + walkWeight G.w (b :: rest.map f)
          have h_ih := ih h_cs_no_none
          -- h_ih: walkWeight G'.w (some b :: rest) = walkWeight G.w ((some b :: rest).map (unsome default))
          -- Compute RHS map:
          have h_map_rest : ((some b :: rest).map (unsome default)) = b :: (rest.map (unsome default)) := by
            simp [unsome]
          rw [h_map_rest] at h_ih
          -- Now h_ih: walkWeight G'.w (some b :: rest) = walkWeight G.w (b :: rest.map (unsome default))
          -- Expand both sides using walkWeight formula
          calc
            walkWeight G.johnsonAugmentedGraph.w (some a :: some b :: rest)
                = G.johnsonAugmentedGraph.w (some a) (some b) +
                  walkWeight G.johnsonAugmentedGraph.w (some b :: rest) := by simp [walkWeight]
            _ = G.w a b + walkWeight G.johnsonAugmentedGraph.w (some b :: rest) := by
              simp [johnsonAugmentedGraph]
            _ = G.w a b + walkWeight G.w (b :: rest.map (unsome default)) := by rw [h_ih]
            _ = walkWeight G.w (a :: b :: rest.map (unsome default)) := by simp [walkWeight]
            _ = walkWeight G.w ((some a :: some b :: rest).map (unsome default)) := by
              unfold unsome; simp

Project a none-free walk in G' to a walk in G with the same weight.


private lemma exists_walk_in_G_of_none_free_walk (x' : V) (c : List (Option V))
    (hc : G.johnsonAugmentedGraph.IsWalkFrom (some x') (some x') c)
    (h_no_none : none ∉ c) :
    ∃ (c' : List V),
      G.IsWalkFrom x' x' c' ∧
      walkWeight G.johnsonAugmentedGraph.w c = walkWeight G.w c' := by
  let c' := c.map (unsome x')
  have h_chain_c' : List.IsChain G.Adj c' :=
    chain_none_free_map G x' c hc.chain h_no_none
  have h_head_c' : c'.head? = some x' := by
    have h_head_c : c.head? = some (some x') := hc.head
    rcases List.head?_eq_some_iff.mp h_head_c with ⟨cs, hc_eq⟩
    subst hc_eq
    simp [c', unsome]
  have h_last_c' : c'.getLast? = some x' := by
    rcases List.getLast?_eq_some_iff.mp hc.last with ⟨q, hq⟩
    subst hq
    simp [c', unsome]
  have h_walk_c' : G.IsWalkFrom x' x' c' :=
    ⟨h_chain_c', h_head_c', h_last_c'⟩
  have h_wt_eq : walkWeight G.johnsonAugmentedGraph.w c = walkWeight G.w c' :=
    walkWeight_none_free_eq G x' c h_no_none
  exact ⟨c', h_walk_c', h_wt_eq⟩

If G has no negative cycles, then johnsonAugmentedGraph also has none.


lemma noNegCycle_johnsonAugmentedGraph (hNC : G.NoNegCycle) :
    G.johnsonAugmentedGraph.NoNegCycle := by
  intro x c hc
  cases x with
  | none =>
    have h_single : c = [none] := walk_from_none_to_none_singleton G c hc
    subst h_single; simp [walkWeight]
  | some x' =>
    have h_head_ne_none : c.head? ≠ some none := by
      rw [hc.head]; simp
    have h_no_none : none ∉ c := chain_no_none G c hc.chain h_head_ne_none
    rcases exists_walk_in_G_of_none_free_walk G x' c hc h_no_none with ⟨c', hc'_walk, hw_eq⟩
    rw [hw_eq]
    exact hNC x' c' hc'_walk

Johnson potential via Bellman-Ford on the augmented graph

A direct walk from none to some v in the augmented graph.


lemma isWalkFrom_none_some (v : V) :
    G.johnsonAugmentedGraph.IsWalkFrom none (some v) [none, some v] := by
  let G' := G.johnsonAugmentedGraph
  have h_chain : List.IsChain G'.Adj [none, some v] := by
    rw [List.isChain_cons]
    refine ⟨?_, List.isChain_singleton _⟩
    intro y hy
    have hy_eq : y = some v := by
      have h_head_val : [some v].head? = some (some v) := by simp
      rw [h_head_val] at hy
      exact (Option.some_inj.mp hy).symm
    subst hy_eq
    unfold G' johnsonAugmentedGraph WeightedGraph.Adj
    simp
  refine ⟨h_chain, ?_, ?_⟩
  · simp
  · simp

The Bellman-Ford distance from none to some v in the augmented graph is finite (not ⊤), because there is a direct zero-weight edge.

Try `simp at h_bound` instead of `simpa using h_bound`

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`
lemma relaxDist_none_some_ne_top (hNC : G.NoNegCycle) (v : V) :
    G.johnsonAugmentedGraph.relaxDist none
      (Fintype.card (Option V) - 1) (some v) ≠ ⊤ := by
  let G' := G.johnsonAugmentedGraph
  have h_no_neg : G'.NoNegCycle := noNegCycle_johnsonAugmentedGraph G hNC
  have h_sd := G'.relaxDist_isShortestDist h_no_neg none (some v)
  rcases h_sd with ⟨h_lower, _⟩
  have h_walk : G'.IsWalkFrom none (some v) [none, some v] :=
    isWalkFrom_none_some G v
  have h_bound : G'.relaxDist none (Fintype.card (Option V) - 1) (some v) ≤
      (walkWeight G'.w [none, some v] : WithTop ℝ) :=
    h_lower [none, some v] h_walk
  have h_wt : (walkWeight G'.w [none, some v] : WithTop ℝ) = (0 : WithTop ℝ) := by
    simp [walkWeight, G', johnsonAugmentedGraph]
  rw [h_wt] at h_bound
  intro h_eq; rw [h_eq] at h_bound; Try `simp at h_bound` instead of `simpa using h_bound`

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`simpa using h_bound

The Johnson potential h(v) is the shortest-path distance from none to some v in the augmented graph, computed by Bellman-Ford.

noncomputable def johnsonPotential (hNC : G.NoNegCycle) (v : V) : ℝ :=
  (G.johnsonAugmentedGraph.relaxDist none
    (Fintype.card (Option V) - 1) (some v)).untop
    (relaxDist_none_some_ne_top G hNC v)

The potential cast to WithTop ℝ equals the Bellman-Ford relaxation.


lemma johnsonPotential_eq (hNC : G.NoNegCycle) (v : V) :
    (G.johnsonPotential hNC v : WithTop ℝ) =
    G.johnsonAugmentedGraph.relaxDist none
      (Fintype.card (Option V) - 1) (some v) := by
  unfold johnsonPotential
  simp [This simp argument is unused:
  relaxDist_none_some_ne_top G hNC v

Hint: Omit it from the simp argument list.
  simp ̵[̵r̵e̵l̵a̵x̵D̵i̵s̵t̵_̵n̵o̵n̵e̵_̵s̵o̵m̵e̵_̵n̵e̵_̵t̵o̵p̵ ̵G̵ ̵h̵N̵C̵ ̵v̵]̵

Note: This linter can be disabled with `set_option linter.unusedSimpArgs false`relaxDist_none_some_ne_top G hNC v]

The Johnson potential is the shortest-path distance from none to some v in the augmented graph.


lemma johnsonPotential_isShortestDist (hNC : G.NoNegCycle) (v : V) :
    G.johnsonAugmentedGraph.IsShortestDist none (some v)
      (G.johnsonPotential hNC v) := by
  rw [johnsonPotential_eq G hNC]
  have h_no_neg : G.johnsonAugmentedGraph.NoNegCycle :=
    noNegCycle_johnsonAugmentedGraph G hNC
  exact G.johnsonAugmentedGraph.relaxDist_isShortestDist h_no_neg none (some v)

Triangle inequality for the Johnson potential

Triangle inequality for the Johnson potential. For every edge (u, v) in G, we have h(v) ≤ h(u) + w(u, v).


lemma johnsonPotential_triangle (hNC : G.NoNegCycle) (u v : V)
    (h_edge : (u, v) ∈ G.edges) :
    G.johnsonPotential hNC v ≤ G.johnsonPotential hNC u + G.w u v := by
  let G' := G.johnsonAugmentedGraph
  let hu := G.johnsonPotential hNC u
  let hv := G.johnsonPotential hNC v
  have h_sd_u : G'.IsShortestDist none (some u) (hu : WithTop ℝ) :=
    johnsonPotential_isShortestDist G hNC u
  have h_sd_v : G'.IsShortestDist none (some v) (hv : WithTop ℝ) :=
    johnsonPotential_isShortestDist G hNC v
  rcases h_sd_u.2 with (h_top | ⟨p_u, hp_walk, hp_weight⟩)
  · have h_fin : (hu : WithTop ℝ) ≠ ⊤ := by simp
    exact absurd h_top h_fin
  · have h_edge' : G'.Adj (some u) (some v) := by
      unfold G' johnsonAugmentedGraph WeightedGraph.Adj; simp [h_edge]
    have h_walk_v : G'.IsWalkFrom none (some v) (p_u ++ [some v]) :=
      IsWalkFrom.append_step (G := G') hp_walk h_edge'
    have h_weight_v : (walkWeight G'.w (p_u ++ [some v]) : WithTop ℝ) =
        (hu : WithTop ℝ) + (G.w u v : WithTop ℝ) := by
      rw [walkWeight_append_step G' hp_walk h_edge']
      rw [hp_weight]
      have h_w_uv : (G'.w (some u) (some v) : WithTop ℝ) = (G.w u v : WithTop ℝ) := by
        unfold G' johnsonAugmentedGraph; simp
      rw [h_w_uv]
    have h_ineq : (hv : WithTop ℝ) ≤ (walkWeight G'.w (p_u ++ [some v]) : WithTop ℝ) :=
      h_sd_v.1 (p_u ++ [some v]) h_walk_v
    rw [h_weight_v] at h_ineq
    exact_mod_cast h_ineq

Nonnegativity of the reweighted graph

With the Johnson potential, every edge weight in the reweighted graph is nonnegative, satisfying Dijkstra's precondition.


lemma reweightedGraph_nonneg (hNC : G.NoNegCycle) :
    (G.reweightedGraph (G.johnsonPotential hNC)).Nonneg := by
  rw [WeightedGraph.Nonneg]
  intro u v h_edge
  rw [w_reweightedGraph, reweightedWeight_eq]
  have h_triangle := johnsonPotential_triangle G hNC u v h_edge
  linarith

With the Johnson potential, the reweighted graph has no negative cycles.

lemma reweightedGraph_noNegCycle (hNC : G.NoNegCycle) :
    (G.reweightedGraph (G.johnsonPotential hNC)).NoNegCycle :=
  noNegCycle_of_nonneg (G := G.reweightedGraph (G.johnsonPotential hNC)) (G.reweightedGraph_nonneg hNC)

Johnson's all-pairs shortest paths

A WithTop ℝ algebra identity for finite adjustments.

private lemma add_sub_add_sub_eq (a : WithTop ℝ) (b c : ℝ) :
    a = (a + (c : WithTop ℝ) - (b : WithTop ℝ)) + (b : WithTop ℝ) - (c : WithTop ℝ) := by
  induction a using WithTop.recTopCoe with
  | top => simp
  | coe a => simp

Johnson's all-pairs shortest-path distance. Run Bellman-Ford from u in the reweighted graph, then adjust by h(v) - h(u).

noncomputable def johnsonDist (hNC : G.NoNegCycle) (u v : V) : WithTop ℝ :=
  let h := G.johnsonPotential hNC
  let d_hat := (G.reweightedGraph h).relaxDist u (Fintype.card V - 1) v
  d_hat + (h v : WithTop ℝ) - (h u : WithTop ℝ)

Theorem (Johnson correctness). johnsonDist hNC u v equals the shortest-path distance δ(u, v) in the original graph G. (CLRS Theorem 23.5)


theorem johnsonDist_isShortestDist (hNC : G.NoNegCycle) (u v : V) :
    G.IsShortestDist u v (G.johnsonDist hNC u v) := by
  let h := G.johnsonPotential hNC
  have h_no_neg : (G.reweightedGraph h).NoNegCycle := G.reweightedGraph_noNegCycle hNC
  let d_hat := (G.reweightedGraph h).relaxDist u (Fintype.card V - 1) v
  have h_sd_dhat : (G.reweightedGraph h).IsShortestDist u v d_hat :=
    (G.reweightedGraph h).relaxDist_isShortestDist h_no_neg u v
  unfold johnsonDist
  -- Let d := d_hat + h(v) - h(u). Then d + h(u) - h(v) = d_hat.
  -- So by reweighted_isShortestDist.mp, G.IsShortestDist u v d.
  have h_algebra : ((d_hat + (h v : WithTop ℝ) - (h u : WithTop ℝ)) +
      (h u : WithTop ℝ) - (h v : WithTop ℝ)) = d_hat :=
    (add_sub_add_sub_eq d_hat (h u) (h v)).symm
  have h_equiv := G.reweighted_isShortestDist h u v
    (d_hat + (h v : WithTop ℝ) - (h u : WithTop ℝ))
  rw [← h_algebra] at h_sd_dhat
  exact h_equiv.mp h_sd_dhat

General triangle inequality

The fundamental triangle inequality for shortest-path distances: for any edge (u, v), the shortest distance to v is at most the shortest distance to u plus the edge weight. No sign assumptions needed.


theorem isShortestDist_edge_ineq (s u v : V) (δ : V → WithTop ℝ)
    (hδ : ∀ t, G.IsShortestDist s t (δ t)) (h_edge : (u, v) ∈ G.edges) :
    δ v ≤ δ u + (G.w u v : WithTop ℝ) := by
  rcases (hδ u).2 with hutop | ⟨q, hq, hqw⟩
  · rw [hutop]; simp
  · have hq_ne : q ≠ [] := hq.ne_nil
    have h_last : q.getLast hq_ne = u := by
      have htemp := List.getLast?_eq_getLast_of_ne_nil hq_ne
      have h_eq_some : some u = some (q.getLast hq_ne) := by
        rw [← hq.last, htemp]
      exact (Option.some_inj.mp h_eq_some).symm
    have h_walk : G.IsWalkFrom s v (q ++ [v]) := by
      refine ⟨?_, ?_, ?_⟩
      · refine hq.chain.append (List.isChain_singleton v) ?_
        intro a ha b hb
        have ha_u : a = u := by
          rw [Option.mem_def, hq.last] at ha
          exact (Option.some.inj ha).symm
        subst ha_u
        have hb_v : b = v := by
          have hsing : [v].head? = some v := by simp
          rw [hsing, Option.mem_def] at hb
          simpa using hb.symm
        subst hb_v
        exact h_edge
      · rw [List.head?_append_of_ne_nil _ hq_ne]
        exact hq.head
      · simp
    have h_weight : (walkWeight G.w (q ++ [v]) : WithTop ℝ) = δ u + (G.w u v : WithTop ℝ) := by
      calc
        (walkWeight G.w (q ++ [v]) : WithTop ℝ) =
            ((walkWeight G.w q + G.w (q.getLast hq_ne) v : ℝ) : WithTop ℝ) := by
          exact_mod_cast walkWeight_append_singleton G.w q hq_ne v
        _ = (walkWeight G.w q : WithTop ℝ) + (G.w (q.getLast hq_ne) v : WithTop ℝ) := by simp
        _ = (walkWeight G.w q : WithTop ℝ) + (G.w u v : WithTop ℝ) := by rw [h_last]
        _ = δ u + (G.w u v : WithTop ℝ) := by rw [hqw]
    have h_bound : δ v ≤ (walkWeight G.w (q ++ [v]) : WithTop ℝ) := (hδ v).1 _ h_walk
    rw [h_weight] at h_bound
    exact h_bound

Every edge in the Johnson-augmented graph targets a some _ vertex.


lemma edge_target_some {u v : Option V}
    (h : (u, v) ∈ G.johnsonAugmentedGraph.edges) : ∃ v', v = some v' := by
  unfold johnsonAugmentedGraph at h
  rcases Finset.mem_union.mp h with (h' | h')
  · rcases Finset.mem_image.mp h' with ⟨v', _, h_eq⟩
    have hv : v = some v' := by
      have h_snd := congrArg Prod.snd h_eq
      simpa using h_snd.symm
    exact ⟨v', hv⟩
  · rcases Finset.mem_image.mp h' with ⟨⟨u', v''⟩, _, h_eq⟩
    have hv : v = some v'' := by
      have h_snd := congrArg Prod.snd h_eq
      simpa using h_snd.symm
    exact ⟨v'', hv⟩

Project a walk in the augmented graph from some a to some b to a walk in G from a to b with the same weight.


lemma walk_johnsonAugmented_some_projection (a b : V) (c : List (Option V))
    (hc : G.johnsonAugmentedGraph.IsWalkFrom (some a) (some b) c) :
    ∃ (c' : List V), G.IsWalkFrom a b c' ∧
      walkWeight G.johnsonAugmentedGraph.w c = walkWeight G.w c' := by
  let G' := G.johnsonAugmentedGraph
  induction c generalizing a b with
  | nil => exact absurd rfl hc.ne_nil
  | cons x xs ih =>
    have hx : x = some a := by
      have hh := hc.head; simp at hh
      -- hh : some x = some (some a)
      simpa using hh
    -- Use rw instead of subst to preserve the IH
    rw [hx] at hc ⊢
    -- Now c = some a :: xs
    rcases (List.isChain_cons_iff G'.Adj (some a) xs).mp hc.chain with
      (hxs_nil | ⟨y, ys, h_adj, h_chain_rest, hxs_eq⟩)
    · -- xs = []
      rw [hxs_nil] at hc ⊢
      -- c = [some a]
      have hab : a = b := by
        have hl := hc.last; simp at hl; simpa using hl
      rw [hab]
      refine ⟨[b], ⟨List.isChain_singleton b, by simp, by simp⟩, ?_⟩
      simp [walkWeight]
    · -- xs = y :: ys; rewrite everything with this
      rw [hxs_eq] at hc ih ⊢
      -- c = some a :: y :: ys
      -- h_adj : G'.Adj (some a) y
      rcases edge_target_some G h_adj with ⟨c', hy⟩
      rw [hy] at hc ih h_adj h_chain_rest ⊢
      -- y = some c'; c = some a :: some c' :: ys
      have h_edge_orig : (a, c') ∈ G.edges :=
        (G.mem_edges_johnsonAugmentedGraph_edge a c').mp h_adj
      -- h_chain_rest : IsChain G'.Adj (some c' :: ys)
      -- Build rest walk
      have h_rest_head : (some c' :: ys).head? = some (some c') := by simp
      have h_rest_last : (some c' :: ys).getLast? = some (some b) := by
        have hl := hc.last; simpa using hl
      have h_rest_walk : G'.IsWalkFrom (some c') (some b) (some c' :: ys) :=
        ⟨h_chain_rest, h_rest_head, h_rest_last⟩
      rcases ih c' b h_rest_walk with ⟨c'_walk, h_walk, h_weight_eq⟩
      -- c'_walk starts at c' and ends at b; deconstruct into c' :: rest
      have h_c'_walk_ne_nil : c'_walk ≠ [] := h_walk.ne_nil
      rcases List.exists_cons_of_ne_nil h_c'_walk_ne_nil with ⟨d, rest, h_c'_walk_eq⟩
      -- d = c' because the walk starts at c'
      have hd_eq_c' : d = c' := by
        have h_walk_head := h_walk.head
        rw [h_c'_walk_eq] at h_walk_head
        simp at h_walk_head
        simpa using h_walk_head
      rw [hd_eq_c'] at h_c'_walk_eq
      -- h_c'_walk_eq : c'_walk = c' :: rest
      have h_walk' : G.IsWalkFrom c' b (c' :: rest) := by
        rwa [h_c'_walk_eq] at h_walk
      have h_weight_eq' : walkWeight G'.w (some c' :: ys) = walkWeight G.w (c' :: rest) := by
        rwa [h_c'_walk_eq] at h_weight_eq
      -- Build full walk a :: c' :: rest
      have h_full_chain : List.IsChain G.Adj (a :: c' :: rest) := by
        rw [List.isChain_cons_iff G.Adj a (c' :: rest)]
        right; exact ⟨c', rest, h_edge_orig, h_walk'.chain, rfl⟩
      have h_full_last : (a :: c' :: rest).getLast? = some b := by
        have hlast := h_walk'.last; simpa using hlast
      have h_full_walk : G.IsWalkFrom a b (a :: c' :: rest) :=
        ⟨h_full_chain, by simp, h_full_last⟩
      have h_weight_eq_full : walkWeight G'.w (some a :: some c' :: ys) =
          walkWeight G.w (a :: c' :: rest) := by
        calc
          walkWeight G'.w (some a :: some c' :: ys) =
              G'.w (some a) (some c') + walkWeight G'.w (some c' :: ys) := rfl
          _ = G.w a c' + walkWeight G'.w (some c' :: ys) := by simp [G', johnsonAugmentedGraph]
          _ = G.w a c' + walkWeight G.w (c' :: rest) := by rw [h_weight_eq']
          _ = walkWeight G.w (a :: c' :: rest) := rfl
      exact ⟨a :: c' :: rest, h_full_walk, h_weight_eq_full⟩

Johnson potential function

The Johnson potential h(v) = δ(none, some v) from Bellman-Ford on the augmented graph. Finite because of the direct none→some v edge (weight 0).

Try `simp at h_final` instead of `simpa using h_final`

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`
lemma johnsonPotential_finite (_hNC : G.NoNegCycle) (v : V) :
    G.johnsonAugmentedGraph.relaxDist none (Fintype.card (Option V) - 1) (some v) ≠ ⊤ := by
  let G' := G.johnsonAugmentedGraph
  have h1 : G'.relaxDist none 1 (some v) ≤ (0 : WithTop ℝ) := by
    calc
      G'.relaxDist none 1 (some v) = G'.relaxStep (G'.relaxDist none 0) (some v) := rfl
      _ ≤ G'.relaxDist none 0 none + (G'.w none (some v) : WithTop ℝ) :=
        G'.relaxStep_le_pred (by simp [G', mem_edges_johnsonAugmentedGraph_source])
      _ = (0 : WithTop ℝ) + (0 : WithTop ℝ) := by simp [G', johnsonAugmentedGraph]
      _ = (0 : WithTop ℝ) := by simp
  have hcardV : 1 ≤ Fintype.card V := Fintype.card_pos_iff.mpr ⟨v⟩
  have hcard_eq : Fintype.card (Option V) - 1 = Fintype.card V := by
    simp [Fintype.card_option]
  rw [hcard_eq]
  have h_noninc : ∀ m, 1 ≤ m → G'.relaxDist none m (some v) ≤ (0 : WithTop ℝ) :=
    Nat.le_induction h1 (fun n hn hle_n =>
      calc
        G'.relaxDist none (n + 1) (some v) ≤ G'.relaxDist none n (some v) :=
          G'.relaxDist_succ_le none n (some v)
        _ ≤ (0 : WithTop ℝ) := hle_n
    )
  have h_final : G'.relaxDist none (Fintype.card V) (some v) ≤ (0 : WithTop ℝ) :=
    h_noninc (Fintype.card V) hcardV
  intro htop; rw [htop] at h_final; Try `simp at h_final` instead of `simpa using h_final`

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`simpa using h_final

Nonnegative reweighted weights



theorem johnsonReweightedNonneg (hNC : G.NoNegCycle) :
    (G.reweightedGraph (G.johnsonPotential hNC)).Nonneg := by
  intro u v h_edge
  have h_edge_orig : (u, v) ∈ G.edges := by
    simpa [reweightedGraph] using h_edge
  have h_triangle := G.johnsonPotential_triangle hNC u v h_edge_orig
  dsimp [reweightedGraph, reweightedWeight]
  linarith

Johnson all-pairs distance

noncomputable def johnsonAllPairsDist (hNC : G.NoNegCycle) (u v : V) : WithTop ℝ :=
  let h := G.johnsonPotential hNC
  let G_hat := G.reweightedGraph h
  let st := G_hat.dijkstraLoop u (Fintype.card V)
  let d_hat := st.d v
  d_hat + (h v : WithTop ℝ) - (h u : WithTop ℝ)

End-to-end correctness



theorem johnsonAllPairsDist_correct (hNC : G.NoNegCycle) (u v : V) :
    G.IsShortestDist u v (G.johnsonAllPairsDist hNC u v) := by
  let h := G.johnsonPotential hNC
  let G_hat := G.reweightedGraph h
  have hnn : G_hat.Nonneg := G.johnsonReweightedNonneg hNC
  have hNC_hat : G_hat.NoNegCycle := G_hat.noNegCycle_of_nonneg hnn
  let δ_hat (v : V) : WithTop ℝ := G_hat.relaxDist u (Fintype.card V - 1) v
  have hδ_hat (v : V) : G_hat.IsShortestDist u v (δ_hat v) :=
    G_hat.relaxDist_isShortestDist hNC_hat u v
  have h_dijkstra (v : V) : (G_hat.dijkstraLoop u (Fintype.card V)).d v = δ_hat v :=
    G_hat.dijkstraLoop_correct hnn u δ_hat hδ_hat (Fintype.card V)
      (le_refl (Fintype.card V)) v
  -- Expand johnsonAllPairsDist using the locally defined h and G_hat
  -- The local h and G_hat are definitionally equal to the ones in the definition
  have h_expand : G.johnsonAllPairsDist hNC u v =
      (G_hat.dijkstraLoop u (Fintype.card V)).d v + (h v : WithTop ℝ) - (h u : WithTop ℝ) := by
    unfold johnsonAllPairsDist
    rfl
  rw [h_expand]
  rw [h_dijkstra v]
  -- Goal: G.IsShortestDist u v (δ_hat v + (h v : WithTop ℝ) - (h u : WithTop ℝ))
  -- reweighted_isShortestDist: (G_hat).IsShortestDist u v (d + h_u - h_v) ↔ G.IsShortestDist u v d
  -- Let d := δ_hat v + h_v - h_u; then d + h_u - h_v = δ_hat v
  have h_fin_hu : (h u : WithTop ℝ) ≠ ⊤ := by simp
  have h_fin_hv : (h v : WithTop ℝ) ≠ ⊤ := by simp
  set d := δ_hat v + (h v : WithTop ℝ) - (h u : WithTop ℝ) with hd
  have h_eq : d + (h u : WithTop ℝ) - (h v : WithTop ℝ) = δ_hat v := by
    dsimp [d]
    rcases Option.ne_none_iff_exists'.mp h_fin_hu with ⟨hu_val, hhu⟩
    rcases Option.ne_none_iff_exists'.mp h_fin_hv with ⟨hv_val, hhv⟩
    rw [hhu, hhv]
    by_cases hδ_top : δ_hat v = ⊤
    · rw [hδ_top]; simp
    · rcases Option.ne_none_iff_exists'.mp hδ_top with ⟨δv_val, hδv⟩
      rw [hδv]; simp
  have h_sd_d : G_hat.IsShortestDist u v (d + (h u : WithTop ℝ) - (h v : WithTop ℝ)) := by
    rw [h_eq]; exact hδ_hat v
  exact (G.reweighted_isShortestDist h u v d).mp h_sd_d

Work bound: O(V² log V + V E log V)

Johnson's algorithm runs one Bellman-Ford pass on the augmented graph (to build the potential h) and then |V| Dijkstra passes on the reweighted graph (which shares G's |E| edges). In the binary-heap cost model of Section 22.3 this is (|V| - 1)·(|V| + |E|) edge relaxations plus |V|·(|V| + |E|)·(log₂|V| + 1) heap operations, i.e. O(V² log V + V E log V).

The two image components of johnsonAugmentedGraph.edges are disjoint: edges out of none never collide with lifted original edges.


lemma johnsonAugmentedGraph_edges_disjoint (G : WeightedGraph V) :
    Disjoint (Finset.image (fun (v : V) => (none, some v)) (Finset.univ : Finset V))
      (Finset.image (fun ((u, v) : V × V) => (some u, some v)) G.edges) := by
  rw [Finset.disjoint_left]
  intro e heA heB
  rcases Finset.mem_image.mp heA with ⟨v, _, rfl⟩
  rcases Finset.mem_image.mp heB with ⟨p, _, h⟩
  have hnone : some p.1 = none := congrArg Prod.fst h
  simp at hnone

The augmented graph has |V| + |E| edges: |V| zero-weight edges out of none plus the |E| original edges.


lemma johnsonAugmentedGraph_edges_card (G : WeightedGraph V) :
    G.johnsonAugmentedGraph.edges.card = Fintype.card V + G.edges.card := by
  have hA_inj : Function.Injective (fun v : V => ((none : Option V), some v)) := by
    intro a b h; exact Option.some.inj (congrArg Prod.snd h)
  have hB_inj : Function.Injective (fun ((u, v) : V × V) => (some u, some v)) := by
    intro a b h
    exact Prod.ext (Option.some.inj (congrArg Prod.fst h)) (Option.some.inj (congrArg Prod.snd h))
  have hA : ((Finset.univ : Finset V).image (fun v : V => ((none : Option V), some v))).card = Fintype.card V :=
    (Finset.card_image_of_injective (Finset.univ : Finset V) hA_inj).trans Finset.card_univ
  have hB : (G.edges.image (fun ((u, v) : V × V) => (some u, some v))).card = G.edges.card :=
    Finset.card_image_of_injective G.edges hB_inj
  unfold johnsonAugmentedGraph
  rw [Finset.card_union_of_disjoint (johnsonAugmentedGraph_edges_disjoint G), hA, hB]

Total Johnson work: one Bellman-Ford pass on the augmented graph (potential) plus |V| Dijkstra passes on the reweighted graph. Noncomputable only because johnsonAugmentedGraph is.

noncomputable def johnsonCost (G : WeightedGraph V) : ℕ :=
  G.johnsonAugmentedGraph.bellmanFordWork +
    Fintype.card V * dijkstraWork (Fintype.card V) G.edges.card

Johnson work bound. In the binary-heap cost model, Johnson's algorithm runs in |V|·(|V| + |E|)·(log₂|V| + 2), i.e. O(V² log V + V E log V).


theorem johnsonCost_eq (G : WeightedGraph V) :
    G.johnsonCost = Fintype.card V * (Fintype.card V + G.edges.card) * (Nat.log2 (Fintype.card V) + 2) := by
  unfold johnsonCost bellmanFordWork dijkstraWork
  rw [johnsonAugmentedGraph_edges_card]
  have hcard : Fintype.card (Option V) - 1 = Fintype.card V := by
    rw [Fintype.card_option]; simp
  rw [hcard]
  ring_nf

O(V² log V + V E log V) work. Since log₂|V| + 2 ≤ 2·(log₂|V| + 1), Johnson's work is at most 2·|V|·(|V| + |E|)·(log₂|V| + 1).


theorem johnsonCost_le (G : WeightedGraph V) :
    G.johnsonCost ≤ 2 * Fintype.card V * (Fintype.card V + G.edges.card) * (Nat.log2 (Fintype.card V) + 1) := by
  rw [johnsonCost_eq]
  calc
    Fintype.card V * (Fintype.card V + G.edges.card) * (Nat.log2 (Fintype.card V) + 2)
        ≤ Fintype.card V * (Fintype.card V + G.edges.card) * (2 * (Nat.log2 (Fintype.card V) + 1)) := by
          exact Nat.mul_le_mul_left _ (by omega)
    _ = 2 * Fintype.card V * (Fintype.card V + G.edges.card) * (Nat.log2 (Fintype.card V) + 1) := by ac_rfl

              end WeightedGraphend Chapter24end CLRS

Definitions and proofs

CLRSLean.FourthEdition.Chapter_23.Section_23_3_Johnsons_Algorithm.Dijkstra

Stored Dijkstra with a list scan queue. The accumulated work counts candidate visits, row-cell evaluations and writes, queue construction, and successful settlements. This indexed real-operation model treats array access and real arithmetic as primitive; it does not claim a binary-heap bound or a bit-complexity bound.

              noncomputable sectionnamespace CLRS.Chapter24.JohnsonExecutionopen MatrixExecution (Stored)structure QueueState (n : Nat) where
  row : Vector
  queue : List (Fin n)def view (st : QueueState n) : WeightedGraph.DijkstraState (Fin n) where
  S := Finset.univ \ st.queue.toFinset
  d := vecRead st.rowdef initState (weights : Stored) (s : Fin n) : QueueState n × Nat :=
  let row := vector (fun v : Fin n => (if v = s then 0 else MatrixExecution.read weights s v, 1))
  let queue := initialQueue s n le_rfl
  (⟨row, queue.1⟩, row.writes + row.visits + queue.2)

@[simp] theorem initState_work (weights : Stored) (s : Fin n) :
    (initState weights s).2 = 3 * n := by
  simp only [initState, vector_writes, initialQueue_visits]
  rw [vector_visits _ 1 (by intros; rfl)]
  omega@[simp] theorem initState_nodup (weights : Stored) (s : Fin n) :
    (initState weights s).1.queue.Nodup := initialQueue_nodup s n le_rfltheorem initState_length_le (weights : Stored) (s : Fin n) :
    (initState weights s).1.queue.length ≤ n := initialQueue_length_le s n le_rfldef MatrixCorrect (G : WeightedGraph (Fin n)) (weights : Stored) : Prop :=
  ∀ i j, MatrixExecution.read weights i j =
    if (i, j) ∈ G.edges then (G.w i j : WithTop ℝ) else ⊤theorem initState_view (G : WeightedGraph (Fin n)) (weights : Stored) (hw : MatrixCorrect G weights)
    (s : Fin n) : view (initState weights s).1 = G.dijkstraInit s := by
  apply WeightedGraph.DijkstraState.ext
  · ext v
    simp [view, initState, initialQueue_mem, v.isLt, WeightedGraph.dijkstraInit]
  · funext v
    simp [view, initState, hw s v, WeightedGraph.dijkstraInit]def relaxRow (weights : Stored) (row : Vector) (u : Fin n) : Vector :=
  vector (fun v : Fin n => (min (vecRead row v) (vecRead row u + MatrixExecution.read weights u v), 1))@[simp] theorem relaxRow_writes (weights : Stored) (row : Vector) (u : Fin n) :
    (relaxRow weights row u).writes = n := vector_writes _@[simp] theorem relaxRow_visits (weights : Stored) (row : Vector) (u : Fin n) :
    (relaxRow weights row u).visits = n := by
  simpa [relaxRow] using vector_visits (fun v : Fin n =>
    (min (vecRead row v) (vecRead row u + MatrixExecution.read weights u v), 1)) 1 (by intros; rfl)def advance (weights : Stored) (st : QueueState n) (u : Fin n) (rest : List (Fin n)) : QueueState n :=
  ⟨relaxRow weights st.row u, rest⟩

theorem advance_view (G : WeightedGraph (Fin n)) (weights : Stored) (hw : MatrixCorrect G weights)
    (st : QueueState n) (hq : st.queue.Nodup) (u : Fin n) (rest : List (Fin n))
    (hp : st.queue.Perm (u :: rest)) :
    view (advance weights st u rest) = G.dijkstraSettle (view st) u := by
  have hnew := List.nodup_cons.mp (hp.nodup_iff.mp hq)
  apply WeightedGraph.DijkstraState.ext
  · ext v
    have hm : v ∈ st.queue ↔ v = u ∨ v ∈ rest := by simpa using hp.mem_iff
    simp only [view, advance, WeightedGraph.dijkstraSettle, Finset.mem_sdiff,
      Finset.mem_univ, List.mem_toFinset, true_and, Finset.mem_insert]
    rw [hm]
    by_cases hv : v = u
    · subst v; simp [hnew.1]
    · simp [hv]
  · funext v
    simp only [view, advance, relaxRow, vector_read, WeightedGraph.dijkstraSettle]
    rw [hw u v]
    by_cases he : (u, v) ∈ G.edges <;> simp [he]def run (weights : Stored) : Nat → QueueState n → QueueState n × Nat
  | 0, st => (st, 0)
  | k + 1, st =>
      let found := extractMin (vecRead st.row) st.queue
      match found.1 with
      | none => (st, found.2)
      | some (u, rest) =>
          let next := advance weights st u rest
          let result := run weights k next
          (result.1, found.2 + next.row.writes + next.row.visits + 1 + result.2)


theorem run_invariant (G : WeightedGraph (Fin n)) (weights : Stored) (hw : MatrixCorrect G weights)
    (hnn : G.Nonneg) (s : Fin n) (δ : Fin n → WithTop ℝ)
    (hδ : ∀ v, G.IsShortestDist s v (δ v)) (k : Nat)
    (st : QueueState n) (hq : st.queue.Nodup)
    (hinv : G.DijkstraInvariant hnn s δ hδ (view st)) :
    G.DijkstraInvariant hnn s δ hδ (view (run weights k st).1) := by
  induction k generalizing st with
  | zero => exact hinv
  | succ k ih =>
    simp only [run]
    split
    · exact hinv
    · rename_i u rest he
      obtain ⟨hp, hm⟩ := extractMin_spec (vecRead st.row) st.queue he
      have hn := List.nodup_cons.mp (hp.nodup_iff.mp hq)
      apply ih _ hn.2
      rw [advance_view G weights hw st hq u rest hp]
      apply G.dijkstraSettle_invariant hnn s δ hδ (view st) hinv u
      · have hu : u ∈ st.queue := hp.mem_iff.mpr (by simp)
        simpa [view] using hu
      · intro v hv
        apply hm v
        simpa [view] using hvtheorem run_queue_empty (weights : Stored) (k : Nat) (st : QueueState n)
    (hlen : st.queue.length ≤ k) : (run weights k st).1.queue = [] := by
  induction k generalizing st with
  | zero => simpa [run] using hlen
  | succ k ih =>
    simp only [run]
    split
    · rename_i he
      exact (extractMin_none (vecRead st.row) st.queue).mp he
    · rename_i u rest he
      obtain ⟨hp, _⟩ := extractMin_spec (vecRead st.row) st.queue he
      have hl := hp.length_eq
      apply ih
      change rest.length ≤ k
      simp only [List.length_cons] at hl
      omega

theorem run_work_le (weights : Stored) (k : Nat) (st : QueueState n)
    (hlen : st.queue.length ≤ n) : (run weights k st).2 ≤ k * (3 * n + 1) := by
  induction k generalizing st with
  | zero => simp [run]
  | succ k ih =>
    simp only [run]
    split
    · simp only [extractMin_visits]
      nlinarith
    · rename_i u rest he
      obtain ⟨hp, _⟩ := extractMin_spec (vecRead st.row) st.queue he
      have hl := hp.length_eq
      have hr : rest.length ≤ n := by simp only [List.length_cons] at hl; omega
      have hind := ih (advance weights st u rest) hr
      simp only [advance] at hind
      simp only [extractMin_visits, advance, relaxRow_writes, relaxRow_visits]
      nlinarithdef dijkstraStored (weights : Stored) (s : Fin n) : Vector × Nat :=
  let initial := initState weights s
  let result := run weights n initial.1
  (result.1.row, initial.2 + result.2)

theorem dijkstraStored_correct (G : WeightedGraph (Fin n)) (weights : Stored)
    (hw : MatrixCorrect G weights) (hnn : G.Nonneg) (s : Fin n) (δ : Fin n → WithTop ℝ)
    (hδ : ∀ v, G.IsShortestDist s v (δ v)) (v : Fin n) :
    vecRead (dijkstraStored weights s).1 v = δ v := by
  have hi : G.DijkstraInvariant hnn s δ hδ (view (initState weights s).1) := by
    rw [initState_view G weights hw s]
    exact G.dijkstraInit_invariant hnn s δ hδ
  have hf := run_invariant G weights hw hnn s δ hδ n (initState weights s).1
    (initState_nodup weights s) hi
  have hq := run_queue_empty weights n (initState weights s).1 (initState_length_le weights s)
  exact hf.hsettled v (by simp [view, hq])theorem dijkstraStored_work_le (weights : Stored) (s : Fin n) :
    (dijkstraStored weights s).2 ≤ 3 * n ^ 2 + 4 * n := by
  have hr := run_work_le weights n (initState weights s).1 (initState_length_le weights s)
  simp only [dijkstraStored, initState_work]
  nlinarithend CLRS.Chapter24.JohnsonExecution

CLRSLean.FourthEdition.Chapter_23.Section_23_3_Johnsons_Algorithm.DijkstraCore

Dijkstra settlement with an explicit minimum

This proof isolates the standard invariant update from the native classical minimum selector. A counted queue may supply any actual unsettled minimum, including its own deterministic tie choice.

namespace CLRS.Chapter24.WeightedGraphvariable {V : Type*} [Fintype V] [DecidableEq V] (G : WeightedGraph V)

Settle the chosen vertex and relax every outgoing edge.

noncomputable def dijkstraSettle (st : DijkstraState V) (u : V) : DijkstraState V :=
  { S := insert u st.S
    d := fun v => if (u, v) ∈ G.edges then min (st.d v) (st.d u + (G.w u v : WithTop ℝ)) else st.d v }

Any unsettled minimum preserves the native Dijkstra invariant.


theorem dijkstraSettle_invariant (hnn : G.Nonneg) (s : V) (δ : V → WithTop ℝ)
    (hδ : ∀ v, G.IsShortestDist s v (δ v))
    (st : DijkstraState V) (h_inv : DijkstraInvariant G hnn s δ hδ st)
    (u : V) (hu_notin_S : u ∉ st.S) (hu_min_all : ∀ y, y ∉ st.S → st.d u ≤ st.d y) :
    DijkstraInvariant G hnn s δ hδ (G.dijkstraSettle st u) := by
  have h_du_eq_δu : st.d u = δ u :=
    G.extractMin_correct_of_invariant hnn s δ hδ st h_inv u hu_notin_S hu_min_all
  let S' := insert u st.S
  let d' := fun v => if (u, v) ∈ G.edges then min (st.d v) (st.d u + (G.w u v : WithTop ℝ)) else st.d v
  have h_s_S' : s ∈ S' := Finset.mem_insert_of_mem h_inv.hsS
  have h_settled' : ∀ x ∈ S', d' x = δ x := by
    intro x hx
    rcases Finset.mem_insert.1 hx with (he | hx_S)
    · subst x
      -- x = u
      dsimp [d']
      by_cases h_edge_uu : (u, u) ∈ G.edges
      · have h_nonneg_w : 0 ≤ G.w u u := hnn u u h_edge_uu
        have h_add : st.d u ≤ st.d u + (G.w u u : WithTop ℝ) := by
          have h_nonneg_w' : (0 : WithTop ℝ) ≤ (G.w u u : WithTop ℝ) := by exact_mod_cast h_nonneg_w
          exact le_add_of_nonneg_right h_nonneg_w'
        simp [h_edge_uu]
        have h_min_eq : min (st.d u) (st.d u + (G.w u u : WithTop ℝ)) = st.d u :=
          min_eq_left h_add
        rw [h_min_eq, h_du_eq_δu]
      · simp [h_edge_uu, h_du_eq_δu]
    · -- x ∈ st.S
      have h_dx_eq_δx : st.d x = δ x := h_inv.hsettled x hx_S
      dsimp [d']
      by_cases h_edge_ux : (u, x) ∈ G.edges
      · have h_ineq : δ x ≤ δ u + (G.w u x : WithTop ℝ) :=
          G.delta_le_delta_add_edge hnn s δ hδ u x h_edge_ux
        have h_add : st.d x ≤ st.d u + (G.w u x : WithTop ℝ) := by
          rw [h_dx_eq_δx, h_du_eq_δu]
          exact h_ineq
        simp [h_edge_ux]
        have h_min_eq : min (st.d x) (st.d u + (G.w u x : WithTop ℝ)) = st.d x :=
          min_eq_left h_add
        rw [h_min_eq, h_dx_eq_δx]
      · simp [h_edge_ux, h_dx_eq_δx]
  have h_htent' : ∀ y ∉ S', ∀ x ∈ S', (x, y) ∈ G.edges → d' y ≤ δ x + (G.w x y : WithTop ℝ) := by
    intro y hy_S' x hx_S' h_edge
    have hy_notin_S : y ∉ st.S := by
      intro hy_S; apply hy_S'; simp [S', hy_S]
    rcases Finset.mem_insert.1 hx_S' with (he | hx_S)
    · subst x
      -- x = u
      dsimp [d']
      have h_edge_uy : (u, y) ∈ G.edges := h_edge
      calc
        (if (u, y) ∈ G.edges then min (st.d y) (st.d u + (G.w u y : WithTop ℝ)) else st.d y)
            = min (st.d y) (st.d u + (G.w u y : WithTop ℝ)) := by simp [h_edge_uy]
        _ ≤ st.d u + (G.w u y : WithTop ℝ) := min_le_right _ _
        _ = δ u + (G.w u y : WithTop ℝ) := by rw [h_du_eq_δu]
    · -- x ∈ st.S
      have h_old_htent : st.d y ≤ δ x + (G.w x y : WithTop ℝ) :=
        h_inv.htent y hy_notin_S x hx_S h_edge
      dsimp [d']
      by_cases h_edge_uy : (u, y) ∈ G.edges
      · calc
          (if (u, y) ∈ G.edges then min (st.d y) (st.d u + (G.w u y : WithTop ℝ)) else st.d y)
              = min (st.d y) (st.d u + (G.w u y : WithTop ℝ)) := by simp [h_edge_uy]
          _ ≤ st.d y := min_le_left _ _
          _ ≤ δ x + (G.w x y : WithTop ℝ) := h_old_htent
      · simp [h_edge_uy, h_old_htent]
  have h_valid' : ∀ y ∉ S', δ y ≤ d' y := by
    intro y hy_S'
    have hy_notin_S : y ∉ st.S := by
      intro hy_S; apply hy_S'; simp [S', hy_S]
    have h_old_valid : δ y ≤ st.d y := h_inv.hvalid y hy_notin_S
    by_cases h_edge_uy : (u, y) ∈ G.edges
    · have h_ineq : δ y ≤ δ u + (G.w u y : WithTop ℝ) :=
        G.delta_le_delta_add_edge hnn s δ hδ u y h_edge_uy
      have h_hvalid_via_add : δ y ≤ st.d u + (G.w u y : WithTop ℝ) := by
        rw [h_du_eq_δu]
        exact h_ineq
      simpa [d', h_edge_uy] using le_min_iff.mpr ⟨h_old_valid, h_hvalid_via_add⟩
    · simpa [d', h_edge_uy] using h_old_valid
  exact ⟨h_s_S', h_settled', h_htent', h_valid'⟩

end CLRS.Chapter24.WeightedGraph

CLRSLean.FourthEdition.Chapter_23.Section_23_3_Johnsons_Algorithm.Execution

Stored Johnson execution with an indexed scan queue

Vertices are explicitly Fin n. The returned array retains one complete row per source. Execution first materializes the augmented edge matrix, computes and caches Bellman–Ford potentials, and materializes one reweighted matrix. Each source then runs stored Dijkstra once and appends its restored row. No shortest-distance oracle is called by the execution definitions.

The returned work is the sum of actual candidate visits, row evaluations/writes, queue construction/settlements, and completed-row appends. In this indexed real-operation model, the dense matrices and scan queue give a cubic upper bound. Indexed array access, graph edge/weight queries, and real arithmetic are primitives; Lean finite-set membership implementation, persistent-array copying, bit complexity, and a binary-heap runtime are outside this model. Correctness requires a graph with no negative cycle; no rejection API is implemented.

noncomputable sectionnamespace CLRS.Chapter24.JohnsonExecution

Execute one source, then materialize the restored original-weight row.

def sourceRow (prepared : Prepared n) (s : Fin n) : Vector × Nat :=
  let computed := dijkstraStored prepared.weights s
  let out := vector (fun v : Fin n =>
    (vecRead computed.1 v + vecRead prepared.potential v - vecRead prepared.potential s, 1))
  (out, computed.2 + out.writes + out.visits)

              

theorem sourceRow_work_le (prepared : Prepared n) (s : Fin n) :
    (sourceRow prepared s).2 ≤ 3 * n ^ 2 + 6 * n := by
  have hd := dijkstraStored_work_le prepared.weights s
  simp only [sourceRow, vector_writes]
  rw [vector_visits _ 1 (by intros; rfl)]
  nlinarith

theorem sourceRow_read (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle) (s v : Fin n) :
    vecRead (sourceRow (prepare G) s).1 v = G.johnsonAllPairsDist hNC s v := by
  let Gh := G.reweightedGraph (G.johnsonPotential hNC)
  have hnn : Gh.Nonneg := G.johnsonReweightedNonneg hNC
  let δ : Fin n → WithTop ℝ := fun v => Gh.relaxDist s (Fintype.card (Fin n) - 1) v
  have hδ : ∀ v, Gh.IsShortestDist s v (δ v) :=
    fun v => Gh.relaxDist_isShortestDist (Gh.noNegCycle_of_nonneg hnn) s v
  have hd := dijkstraStored_correct Gh (prepare G).weights (prepare_weights G hNC) hnn s δ hδ v
  have hl := Gh.dijkstraLoop_correct hnn s δ hδ n (by simp) v
  simp only [sourceRow, vector_read, prepare_potential G hNC, hd]
  simpa [WeightedGraph.johnsonAllPairsDist, Gh] using
    congrArg (fun d : WithTop ℝ => d + (G.johnsonPotential hNC v : WithTop ℝ) -
      (G.johnsonPotential hNC s : WithTop ℝ)) hl.symmstructure Result (n : Nat) where
  rows : Array Vector
  work : Nat

Append one completed source row per iteration; the inner cell loop never reruns Dijkstra.

def sourceRows (prepared : Prepared n) : (k : Nat) → k ≤ n → Result n
  | 0, _ => ⟨#[], 0⟩
  | k + 1, hk =>
    let previous := sourceRows prepared k (by omega)
    let next := sourceRow prepared ⟨k, by omega⟩
    ⟨previous.rows.push next.1, previous.work + next.2 + 1⟩

              @[simp] theorem sourceRows_size (prepared : Prepared n) (k : Nat) (hk : k ≤ n) :
    (sourceRows prepared k hk).rows.size = k := by
  induction k with
  | zero => rfl
  | succ k ih => simp [sourceRows, ih]

theorem sourceRows_get (prepared : Prepared n) (k : Nat) (hk : k ≤ n)
    (i : Nat) (hi : i < k) :
    (sourceRows prepared k hk).rows[i]? = some (sourceRow prepared ⟨i, by omega⟩).1 := by
  induction k with
  | zero => omega
  | succ k ih =>
    by_cases he : i = k
    · subst i
      simp only [sourceRows, Array.getElem?_push, sourceRows_size, ite_true]
    · have hi' : i < k := by omega
      simp only [sourceRows, Array.getElem?_push]
      rw [sourceRows_size]
      simp only [he, if_false]
      exact ih (by omega) hi'

theorem sourceRows_work_le (prepared : Prepared n) (k : Nat) (hk : k ≤ n) :
    (sourceRows prepared k hk).work ≤ k * (3 * n ^ 2 + 6 * n + 1) := by
  induction k with
  | zero => simp [sourceRows]
  | succ k ih =>
    have hs := sourceRow_work_le prepared (⟨k, by omega⟩ : Fin n)
    have hp := ih (by omega)
    simp only [sourceRows]
    nlinarithdef read (result : Result n) (i j : Fin n) : WithTop ℝ :=
  ((result.rows[i.val]?).map (fun row => vecRead row j)).getD ⊤

Stored Johnson execution for explicitly indexed vertices, with a scan priority queue.

def johnsonStored (G : WeightedGraph (Fin n)) : Result n :=
  let prepared := prepare G
  let rows := sourceRows prepared n (Nat.le_refl n)
  { rows with work := prepared.work + rows.work }

              @[simp] theorem johnsonStored_size (G : WeightedGraph (Fin n)) :
    (johnsonStored G).rows.size = n := sourceRows_size _ _ _theorem johnsonStored_read (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle) (i j : Fin n) :
    read (johnsonStored G) i j = G.johnsonAllPairsDist hNC i j := by
  simp only [read, johnsonStored, sourceRows_get _ n (Nat.le_refl n) i.val i.isLt,
    Option.map_some, Option.getD_some]
  exact sourceRow_read G hNC i j

theorem johnsonStored_correct (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle) (i j : Fin n) :
    G.IsShortestDist i j (read (johnsonStored G) i j) := by
  rw [johnsonStored_read G hNC]
  exact G.johnsonAllPairsDist_correct hNC i j

Every source has its own retained, fully materialized row.

theorem johnsonStored_row (G : WeightedGraph (Fin n)) (i : Fin n) :
    (johnsonStored G).rows[i.val]? = some (sourceRow (prepare G) i).1 :=
  sourceRows_get _ n (Nat.le_refl n) i.val i.isLt

@[simp] theorem sourceRow_width (prepared : Prepared n) (s : Fin n) :
    (sourceRow prepared s).1.cells.size = n := by
  simp [sourceRow, vector]

A specification bridge useful for finite examples, independent of Dijkstra tie choices.


theorem johnsonStored_eq_relaxDist (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle)
    (i j : Fin n) : read (johnsonStored G) i j = G.relaxDist i (n - 1) j := by
  have hs := johnsonStored_correct G hNC i j
  have hr : G.IsShortestDist i j (G.relaxDist i (n - 1) j) := by
    simpa using G.relaxDist_isShortestDist hNC i j
  have hle {a b : WithTop ℝ} (ha : G.IsShortestDist i j a)
      (hb : G.IsShortestDist i j b) : a ≤ b := by
    obtain ht | ⟨p, hp, hw⟩ := hb.2
    · rw [ht]; exact le_top
    · rw [← hw]; exact ha.1 p hp
  exact le_antisymm (hle hs hr) (hle hr hs)

              theorem johnsonStored_work_le_polynomial (G : WeightedGraph (Fin n)) :
    (johnsonStored G).work ≤ 4 * n ^ 3 + 14 * n ^ 2 + 12 * n + 4 := by
  have hs := sourceRows_work_le (prepare G) n (Nat.le_refl n)
  simp only [johnsonStored, prepare_work]
  nlinariththeorem johnsonStored_work_le (G : WeightedGraph (Fin n)) :
    (johnsonStored G).work ≤ 16 * (n + 1) ^ 3 := by
  have hs := johnsonStored_work_le_polynomial G
  nlinarith [Nat.zero_le (n ^ 3), Nat.zero_le (n ^ 2)]end CLRS.Chapter24.JohnsonExecution

CLRSLean.FourthEdition.Chapter_23.Section_23_3_Johnsons_Algorithm.Potential

Stored Johnson preprocessing: materialize the augmented edge matrix, execute synchronous Bellman–Ford rounds, retain the potential vector, and materialize one reweighted matrix. Correctness assumes absence of negative cycles; execution does not implement a negative-cycle failure API. Counters use the indexed real-operation model of the storage primitives.

              noncomputable sectionnamespace CLRS.Chapter24.JohnsonExecutionopen MatrixExecution (Stored)def potentialPhase (G : WeightedGraph (Fin n)) : Vector × Nat :=
  let e := finSuccEquiv n
  let weights := edgeTable G.johnsonAugmentedGraph e
  let relaxed := bfLoop weights (e.symm none) n
  let potential := vector (fun v : Fin n => (vecRead relaxed.1 (e.symm (some v)), 1))
  (potential, weights.cellWrites + weights.candidateVisits + relaxed.2 + potential.writes + potential.visits)


theorem potentialPhase_read (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle) (v : Fin n) :
    vecRead (potentialPhase G).1 v = (G.johnsonPotential hNC v : WithTop ℝ) := by
  unfold potentialPhase
  rw [vector_read, bfLoop_read G.johnsonAugmentedGraph (finSuccEquiv n) _
    (edgeTable_read G.johnsonAugmentedGraph (finSuccEquiv n))]
  simpa using (G.johnsonPotential_eq hNC v).symm


theorem potentialPhase_work (G : WeightedGraph (Fin n)) :
    (potentialPhase G).2 = 2 * (n + 1) ^ 2 + 2 * (n + 1) +
      n * ((n + 1) * (n + 3)) + 2 * n := by
  simp only [potentialPhase, edgeTable_writes, edgeTable_visits, bfLoop_work, vector_writes]
  rw [vector_visits _ 1 (by intros; rfl)]
  ringdef reweightTable (G : WeightedGraph (Fin n)) (potential : Vector) : Stored :=
  MatrixExecution.tabulate (fun u v : Fin n =>
    (if (u, v) ∈ G.edges then (G.w u v : WithTop ℝ) + vecRead potential u - vecRead potential v else ⊤, 1))@[simp] theorem reweightTable_writes (G : WeightedGraph (Fin n)) (potential : Vector) :
    (reweightTable G potential).cellWrites = n * n := MatrixExecution.tabulate_writes _@[simp] theorem reweightTable_visits (G : WeightedGraph (Fin n)) (potential : Vector) :
    (reweightTable G potential).candidateVisits = n * n := by
  unfold reweightTable
  simpa using MatrixExecution.tabulate_visits
    (fun u v : Fin n =>
      (if (u, v) ∈ G.edges then (G.w u v : WithTop ℝ) + vecRead potential u - vecRead potential v else ⊤, 1))
    1 (by intros; rfl)structure Prepared (n : Nat) where
  potential : Vector
  weights : Stored
  work : Natdef prepare (G : WeightedGraph (Fin n)) : Prepared n :=
  let potential := potentialPhase G
  let weights := reweightTable G potential.1
  ⟨potential.1, weights, potential.2 + weights.cellWrites + weights.candidateVisits⟩@[simp] theorem prepare_potential (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle) (v : Fin n) :
    vecRead (prepare G).potential v = (G.johnsonPotential hNC v : WithTop ℝ) :=
  potentialPhase_read G hNC vtheorem prepare_weights (G : WeightedGraph (Fin n)) (hNC : G.NoNegCycle) (u v : Fin n) :
    MatrixExecution.read (prepare G).weights u v =
      if (u, v) ∈ (G.reweightedGraph (G.johnsonPotential hNC)).edges then
        ((G.reweightedGraph (G.johnsonPotential hNC)).w u v : WithTop ℝ) else ⊤ := by
  simp only [prepare, reweightTable, MatrixExecution.tabulate_read, potentialPhase_read G hNC]
  simp only [WeightedGraph.reweightedGraph, WeightedGraph.reweightedWeight]
  split <;> simp [WithTop.coe_add]theorem prepare_work (G : WeightedGraph (Fin n)) :
    (prepare G).work = n ^ 3 + 8 * n ^ 2 + 11 * n + 4 := by
  simp only [prepare, potentialPhase_work, reweightTable_writes, reweightTable_visits]
  ringend CLRS.Chapter24.JohnsonExecution

CLRSLean.FourthEdition.Chapter_23.Section_23_3_Johnsons_Algorithm.Queue

Counted scan queue

Minimum extraction traverses an explicit queue and returns the remaining queue without its selected minimum. Its counter counts actual candidate visits. This is a linear scan queue, with no binary-heap logarithmic bound claimed.

noncomputable sectionnamespace CLRS.Chapter24.JohnsonExecution

Scan a list queue, returning its minimum and the remaining queue. Each nonempty recursive frame records one actual candidate visit.

def extractMin (d : α → WithTop ℝ) : List α → Option (α × List α) × Nat
  | [] => (none, 0)
  | x :: xs =>
      let tail := extractMin d xs
      match tail.1 with
      | none => (some (x, []), tail.2 + 1)
      | some (y, ys) =>
          if d x ≤ d y then (some (x, xs), tail.2 + 1)
          else (some (y, x :: ys), tail.2 + 1)

              @[simp] theorem extractMin_visits (d : α → WithTop ℝ) (xs : List α) :
    (extractMin d xs).2 = xs.length := by
  induction xs with
  | nil => rfl
  | cons x xs ih =>
    simp only [extractMin]
    split
    · simp [ih]
    · split <;> simp [ih]@[simp] theorem extractMin_none (d : α → WithTop ℝ) (xs : List α) :
    (extractMin d xs).1 = none ↔ xs = [] := by
  cases xs with
  | nil => simp [extractMin]
  | cons x xs =>
    simp only [extractMin]
    split
    · simp
    · split <;> simptheorem extractMin_spec (d : α → WithTop ℝ) (xs : List α) {u rest}
    (h : (extractMin d xs).1 = some (u, rest)) :
    xs.Perm (u :: rest) ∧ ∀ v ∈ xs, d u ≤ d v := by
  induction xs generalizing u rest with
  | nil => simp [extractMin] at h
  | cons x xs ih =>
    simp only [extractMin] at h
    split at h
    · rename_i ht
      have he := (extractMin_none d xs).mp ht
      subst xs
      have he : (x, []) = (u, rest) := Option.some.inj h
      cases he
      exact ⟨.refl _, by simp⟩
    · rename_i y ys ht
      obtain ⟨hp, hm⟩ := ih ht
      split at h
      · rename_i hxy
        have he : (x, xs) = (u, rest) := Option.some.inj h
        cases he
        refine ⟨.refl _, ?_⟩
        intro v hv
        rcases List.mem_cons.mp hv with rfl | hv
        · exact le_rfl
        · exact le_trans hxy (hm v hv)
      · rename_i hxy
        have he : (y, x :: ys) = (u, rest) := Option.some.inj h
        cases he
        refine ⟨(List.Perm.cons x hp).trans (.swap u x ys), ?_⟩
        intro v hv
        rcases List.mem_cons.mp hv with rfl | hv
        · exact le_of_lt (lt_of_not_ge hxy)
        · exact hm v hv

Build the explicit unsettled queue, counting each enumerated vertex.

def initialQueue (s : Fin n) : (k : Nat) → k ≤ n → List (Fin n) × Nat
  | 0, _ => ([], 0)
  | k + 1, hk =>
      let rest := initialQueue s k (by omega)
      let v : Fin n := ⟨k, by omega⟩
      (if v = s then rest.1 else v :: rest.1, rest.2 + 1)

              @[simp] theorem initialQueue_visits (s : Fin n) (k : Nat) (hk : k ≤ n) :
    (initialQueue s k hk).2 = k := by
  induction k with
  | zero => rfl
  | succ k ih => simp [initialQueue, ih]

theorem initialQueue_mem (s : Fin n) (k : Nat) (hk : k ≤ n) (v : Fin n) :
    v ∈ (initialQueue s k hk).1 ↔ v.val < k ∧ v ≠ s := by
  induction k with
  | zero => simp [initialQueue]
  | succ k ih =>
    simp only [initialQueue]
    split
    · rename_i he
      rw [ih]
      constructor
      · intro h; exact ⟨by omega, h.2⟩
      · rintro ⟨hv, hvs⟩
        refine ⟨?_, hvs⟩
        have hne : v.val ≠ k := by
          intro heq
          apply hvs
          exact (Fin.ext heq).trans he
        omega
    · rename_i he
      rw [List.mem_cons, ih]
      constructor
      · rintro (rfl | ⟨hv, hvs⟩)
        · exact ⟨by simp, he⟩
        · exact ⟨by omega, hvs⟩
      · rintro ⟨hv, hvs⟩
        by_cases heq : v.val = k
        · left; exact Fin.ext heq
        · right; exact ⟨by omega, hvs⟩

theorem initialQueue_nodup (s : Fin n) (k : Nat) (hk : k ≤ n) :
    (initialQueue s k hk).1.Nodup := by
  induction k with
  | zero => simp [initialQueue]
  | succ k ih =>
    simp only [initialQueue]
    split
    · exact ih _
    · rw [List.nodup_cons]
      exact ⟨by simp [initialQueue_mem], ih _⟩

theorem initialQueue_length_le (s : Fin n) (k : Nat) (hk : k ≤ n) :
    (initialQueue s k hk).1.length ≤ k := by
  induction k with
  | zero => simp [initialQueue]
  | succ k ih =>
    have hh := ih (by omega)
    simp only [initialQueue]
    split
    · exact Nat.le_trans hh (Nat.le_succ _)
    · simpa only [List.length_cons] using Nat.succ_le_succ hhend CLRS.Chapter24.JohnsonExecution

CLRSLean.FourthEdition.Chapter_23.Section_23_3_Johnsons_Algorithm.Storage

Stored rows and synchronous Bellman–Ford execution

Rows and edge tables are materialized before use. Every cell append and scanned candidate contributes to a counter returned by the actual loops. Bellman–Ford reads the previous stored row and a cached dense edge table; no recursive relaxation oracle is called during a round. Exact real arithmetic, graph edge/weight queries, and indexed array access are abstract primitives. The model excludes Lean finite-set membership implementation, bit costs, and persistent-array copying.

              noncomputable sectionnamespace CLRS.Chapter24.JohnsonExecutionopen CLRS.Chapter15.DPExecutionopen MatrixExecution (Stored)abbrev Vector := Row (WithTop ℝ)def vecRead (row : Vector) (i : Fin n) : WithTop ℝ := row.cells[i.val]?.getD ⊤def vector (f : Fin n → WithTop ℝ × Nat) : Vector :=
  buildRow (fun i => if hi : i < n then f ⟨i, hi⟩ else (⊤, 0)) n@[simp] theorem vector_read (f : Fin n → WithTop ℝ × Nat) (i : Fin n) :
    vecRead (vector f) i = (f i).1 := by
  simp [vecRead, vector, i.isLt]@[simp] theorem vector_writes (f : Fin n → WithTop ℝ × Nat) : (vector f).writes = n := by
  simp [vector]

theorem vector_visits (f : Fin n → WithTop ℝ × Nat) (c : Nat) (hc : ∀ i, (f i).2 = c) :
    (vector f).visits = n * c := by
  rw [vector, buildRow_visits]
  calc
    _ = ∑ _i ∈ Finset.range n, c := by
      apply Finset.sum_congr rfl
      intro i hi
      simp [Finset.mem_range.mp hi, hc]
    _ = _ := by simpvariable {V : Type} [Fintype V] [DecidableEq V]def edgeTable (G : WeightedGraph V) (e : Fin n ≃ V) : Stored :=
  MatrixExecution.tabulate (fun i j => (if (e i, e j) ∈ G.edges then (G.w (e i) (e j) : WithTop ℝ) else ⊤, 1))@[simp] theorem edgeTable_read (G : WeightedGraph V) (e : Fin n ≃ V) (i j : Fin n) :
    MatrixExecution.read (edgeTable G e) i j =
      if (e i, e j) ∈ G.edges then (G.w (e i) (e j) : WithTop ℝ) else ⊤ := by
  simp [edgeTable]@[simp] theorem edgeTable_writes (G : WeightedGraph V) (e : Fin n ≃ V) :
    (edgeTable G e).cellWrites = n * n := MatrixExecution.tabulate_writes _@[simp] theorem edgeTable_visits (G : WeightedGraph V) (e : Fin n ≃ V) :
    (edgeTable G e).candidateVisits = n * n := by
  simpa [edgeTable] using MatrixExecution.tabulate_visits
    (fun i j : Fin n => (if (e i, e j) ∈ G.edges then (G.w (e i) (e j) : WithTop ℝ) else ⊤, 1)) 1
    (by intros; rfl)def bfRound (weights : Stored) (prev : Vector) (n : Nat) : Vector :=
  vector (fun v : Fin n =>
    let candidates := MatrixExecution.scanFin (fun u : Fin n =>
      vecRead prev u + MatrixExecution.read weights u v)
    (min (vecRead prev v) candidates.1, candidates.2 + 1))@[simp] theorem bfRound_writes (weights : Stored) (prev : Vector) (n : Nat) :
    (bfRound weights prev n).writes = n := vector_writes _@[simp] theorem bfRound_visits (weights : Stored) (prev : Vector) (n : Nat) :
    (bfRound weights prev n).visits = n * (n + 1) := by
  apply vector_visits
  intro i
  simpomit [DecidableEq V] in
private theorem inf_fin_equiv (e : Fin n ≃ V) (f : V → WithTop ℝ) :
    Finset.univ.inf (fun i : Fin n => f (e i)) = Finset.univ.inf f := by
  apply le_antisymm
  · apply Finset.le_inf
    intro v _
    simpa using (Finset.inf_le (s := Finset.univ) (f := fun i : Fin n => f (e i))
      (Finset.mem_univ (e.symm v)))
  · apply Finset.le_inf
    intro i _
    exact Finset.inf_le (Finset.mem_univ (e i))

theorem bfRound_read (G : WeightedGraph V) (e : Fin n ≃ V) (weights : Stored)
    (hw : ∀ i j, MatrixExecution.read weights i j =
      if (e i, e j) ∈ G.edges then (G.w (e i) (e j) : WithTop ℝ) else ⊤)
    (prev : Vector) (d : V → WithTop ℝ) (hd : ∀ i, vecRead prev i = d (e i)) (v : Fin n) :
    vecRead (bfRound weights prev n) v = G.relaxStep d (e v) := by
  simp only [bfRound, vector_read, MatrixExecution.scanFin_value, hd, hw]
  simp_rw [add_ite, add_top]
  rw [inf_fin_equiv e (fun u => if (u, e v) ∈ G.edges then d u + (G.w u (e v) : WithTop ℝ) else ⊤)]
  simp [WeightedGraph.relaxStep, WeightedGraph.preds, Finset.inf_ite]def bfLoop (weights : Stored) (s : Fin n) : Nat → Vector × Nat
  | 0 =>
      let initial := vector (fun i : Fin n => (if i = s then 0 else ⊤, 1))
      (initial, initial.writes + initial.visits)
  | k + 1 =>
      let prev := bfLoop weights s k
      let next := bfRound weights prev.1 n
      (next, prev.2 + next.writes + next.visits)

theorem bfLoop_read (G : WeightedGraph V) (e : Fin n ≃ V) (weights : Stored)
    (hw : ∀ i j, MatrixExecution.read weights i j =
      if (e i, e j) ∈ G.edges then (G.w (e i) (e j) : WithTop ℝ) else ⊤)
    (s : Fin n) (k : Nat) (v : Fin n) :
    vecRead (bfLoop weights s k).1 v = G.relaxDist (e s) k (e v) := by
  induction k generalizing v with
  | zero => simp [bfLoop, WeightedGraph.relaxDist, e.injective.eq_iff]
  | succ k ih =>
    change vecRead (bfRound weights (bfLoop weights s k).1 n) v = _
    rw [bfRound_read G e weights hw _ (G.relaxDist (e s) k) (fun i => ?_)]
    · rfl
    · exact ih i

@[simp] theorem bfLoop_work (weights : Stored) (s : Fin n) (k : Nat) :
    (bfLoop weights s k).2 = 2 * n + k * (n * (n + 2)) := by
  induction k with
  | zero =>
    simp only [bfLoop, vector_writes]
    rw [vector_visits _ 1 (by intros; rfl)]
    omega
  | succ k ih =>
    simp only [bfLoop, bfRound_writes, bfRound_visits, ih]
    ringend CLRS.Chapter24.JohnsonExecution