Problem: k_smallest

Your task is to implement the Quick Select algorithm to find the k-th smallest element and verify its correctness in Lean. Given a list of integers and an index k, return the k-th smallest element (0-indexed). Postconditions: The result is the element that would be at index k if the array were sorted. Preconditions: k < array.size. Implementation: Recursive partition-based selection (like Quick Sort partition). Return the element if the pivot index is k, otherwise recurse into the appropriate sub-array. Time complexity should be O(N) on average.


import Mathlib

-- Precondition definitions
@[reducible, simp]
def quick_select_precond (v : Array Int) (k : Nat) : Prop :=
  -- !benchmark @start precond
  k < v.size
  -- !benchmark @end precond

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

-- Main function definitions
def quick_select (v : Array Int) (k : Nat) (h_precond : quick_select_precond v k) : Int :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition definitions
@[reducible, simp]
def quick_select_postcond (v : Array Int) (k : Nat) (result : Int) (h_precond : quick_select_precond v k) : Prop :=
  -- !benchmark @start postcond
  ∃ sorted : List Int,
    sorted.length = v.size ∧
    List.Sorted (· ≤ ·) sorted ∧
    sorted.getD k 0 = result ∧
    sorted.Perm v.toList
  -- !benchmark @end postcond

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

-- Proof content
theorem quick_select_postcond_satisfied (v : Array Int) (k : Nat) (h_precond : quick_select_precond v k) :
    quick_select_postcond v k (quick_select v k h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof