Problem: kruskal

Your task is to implement Kruskal's Algorithm to find the Minimum Spanning Tree (MST) of a connected, undirected, weighted graph and verify its functional correctness in Lean using Mathlib. The graph is represented as an adjacency list with Int weights. Preconditions: You can assume the graph is connected (a path exists between any two nodes). You can assume the graph contains at most 100,000 nodes and edge weights are between -100,000 and 100,000 to prevent integer overflow. Crucially, the graph is a Simple Graph (no parallel edges between the same two nodes). Requirements: Implement the function kruskal_mst which returns a list of edges (source, target, weight). (1) The function must sort all edges by weight and iterate through them. (2) It should add an edge to the result only if it connects two previously disconnected components (i.e., does not form a cycle). (3) Global Optimality: You must prove that the sum of weights of your returned edges is less than or equal to the sum of weights of any other possible spanning tree of the graph. Verification Challenges: (1) Union-Find (Disjoint Set): To efficiently implement Kruskal's, you typically use a Union-Find data structure; you must verify that your find and union operations correctly track component connectivity. (2) Cycle Prevention: Proving that your check (typically find(u) != find(v)) correctly predicts whether adding edge $(u, v)$ would create a cycle. (3) Forest Invariant: Unlike Prim's (which grows one tree), Kruskal's grows a forest of trees; you must prove the invariant that this forest eventually merges into a single spanning tree and that every edge added is a "safe edge" for some MST.

import Mathlib

structure WeightedGraph where
  adj : Array (Array (Nat × Int))

def WeightedGraph.size (g : WeightedGraph) : Nat :=
  g.adj.size

def WeightedGraph.has_edge (g : WeightedGraph) (u v : Nat) (w : Int) : Prop :=
  u < g.size ∧
  ∃ pair, pair ∈ g.adj.getD u #[] ∧ pair.1 = v ∧ pair.2 = w

def WeightedGraph.well_formed (g : WeightedGraph) : Prop :=
  ∀ u, u < g.size →
    -- Bounds check
    (∀ pair, pair ∈ g.adj.getD u #[] → pair.1 < g.size) ∧
    -- Simple graph check (matches Dafny)
    (∀ p1 p2, p1 ∈ g.adj.getD u #[] → p2 ∈ g.adj.getD u #[] → p1.1 = p2.1 → p1 = p2)

structure MSTEdge where
  u : Nat
  v : Nat
  w : Int

-- 1. Validity
def MSTEdge.is_valid (e : MSTEdge) (g : WeightedGraph) : Prop :=
  g.has_edge e.u e.v e.w

-- 2. Connectivity for Solution Edges (Undirected View)
-- Dafny: `edge_in_set` checks (u, v) || (v, u)
def edges_contain (edges : List MSTEdge) (u v : Nat) : Prop :=
  ∃ e, e ∈ edges ∧ ((e.u = u ∧ e.v = v) ∨ (e.u = v ∧ e.v = u))

def path_follows_edges (edges : List MSTEdge) (p : List Nat) : Prop :=
  p.length > 0 ∧
  ∀ i, i + 1 < p.length →
    edges_contain edges (p.getD i 0) (p.getD (i + 1) 0)

def is_connected_with_edges (g : WeightedGraph) (edges : List MSTEdge) : Prop :=
  ∀ u v, u < g.size → v < g.size →
    ∃ p, p.head? = some u ∧ p.getLast? = some v ∧ path_follows_edges edges p

-- 3. Spanning Tree Definition
def is_spanning_tree (g : WeightedGraph) (edges : List MSTEdge) : Prop :=
  edges.length = g.size - 1 ∧
  (∀ e, e ∈ edges → e.is_valid g) ∧
  is_connected_with_edges g edges
  -- Note: |V|-1 edges + Connected implies Acyclic (Tree)

def total_weight (edges : List MSTEdge) : Int :=
  edges.foldl (λ sum e => sum + e.w) 0

-- 4. Optimality (MST)
def is_mst (g : WeightedGraph) (edges : List MSTEdge) : Prop :=
  is_spanning_tree g edges ∧
  ∀ other, is_spanning_tree g other → total_weight edges ≤ total_weight other

-- 5. Graph Connectivity (Precondition)
-- Dafny uses directed edges for the graph path check: `graph_has_any_edge`
def path_follows_graph (g : WeightedGraph) (p : List Nat) : Prop :=
  p.length > 0 ∧
  ∀ i, i + 1 < p.length →
    ∃ w, g.has_edge (p.getD i 0) (p.getD (i + 1) 0) w

def WeightedGraph.is_connected (g : WeightedGraph) : Prop :=
  ∀ u v, u < g.size → v < g.size →
    ∃ p, p.head? = some u ∧ p.getLast? = some v ∧ path_follows_graph g p

def WeightedGraph.weights_bounded (g : WeightedGraph) : Prop :=
  g.size ≤ 100000 ∧
  ∀ u v w, g.has_edge u v w → w ≥ -100000 ∧ w ≤ 100000

-- Precondition definitions
@[reducible, simp]
def kruskal_mst_precond (graph : WeightedGraph) : Prop :=
  -- !benchmark @start precond
  graph.well_formed ∧
  graph.size > 0 ∧
  graph.is_connected ∧
  graph.weights_bounded
  -- !benchmark @end precond

-- !benchmark @start auxcode
-- !benchmark @end auxcode

-- Main function definition
def kruskal_mst (graph : WeightedGraph)
    (h_precond : kruskal_mst_precond graph) : List MSTEdge :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition definitions
@[reducible, simp]
def kruskal_mst_postcond (graph : WeightedGraph) (result : List MSTEdge)
    (_ : kruskal_mst_precond graph) : Prop :=
  -- !benchmark @start postcond
  is_mst graph result
  -- !benchmark @end postcond

-- !benchmark @start lemma
-- !benchmark @end lemma

-- Proof content
theorem kruskal_mst_postcond_satisfied (graph : WeightedGraph)
    (h_precond : kruskal_mst_precond graph) :
    kruskal_mst_postcond graph (kruskal_mst graph h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof