-----Description-----
This task requires writing a Lean 4 method that returns the k most frequent elements from a list of integers. The method should count the frequency of each distinct element in the list and return the k elements with the highest frequency.

-----Input-----
The input consists of two values:
nums: A list of integers, possibly with duplicates.
k: A natural number indicating how many of the most frequent elements to return. Assuming k <= # of distinct elements in nums.

-----Output-----
The output is a list of integers:
Returns exactly k integers representing the elements that appear most frequently in the input list in order from higher frequency to lower frequency.
If two numbers have the same frequency, any consistent ordering is acceptable (no strict tie-breaking requirement).

Difficulty: advanced

Lean 4 solution:
-- !benchmark @start import type=solution

-- !benchmark @end import

-- !benchmark @start solution_aux

-- !benchmark @end solution_aux

-- !benchmark @start precond_aux

-- !benchmark @end precond_aux
@[reducible]
def topKFrequent_precond (nums : List Int) (k : Nat) : Prop :=
  -- !benchmark @start precond
  k ≤ nums.eraseDups.length
  -- !benchmark @end precond


-- !benchmark @start code_aux

-- !benchmark @end code_aux


def topKFrequent (nums : List Int) (k : Nat) (h_precond : topKFrequent_precond (nums) (k)) : List Int :=
  -- !benchmark @start code
  -- Build frequency list maintaining first occurrence order (deterministic)
  let freqList : List (Int × Nat) :=
    nums.foldl (init := []) fun acc n =>
      match acc.find? (fun (key, _) => key == n) with
      | some _ => acc.map (fun (key, cnt) => if key == n then (key, cnt + 1) else (key, cnt))
      | none => acc ++ [(n, 1)]
  -- Sort by frequency (descending), stable sort preserves first-occurrence order for ties
  let sorted := freqList.foldl
    (fun acc pair =>
      let (x, cx) := pair
      let rec insertSorted (xs : List (Int × Nat)) : List (Int × Nat) :=
        match xs with
        | [] => [pair]
        | (y, cy) :: ys =>
          if cx > cy then
            pair :: (y, cy) :: ys
          else
            (y, cy) :: insertSorted ys
      insertSorted acc
    ) []

  sorted.take k |>.map (fun (n, _) => n)
  -- !benchmark @end code


-- !benchmark @start postcond_aux

-- !benchmark @end postcond_aux


@[reducible]
def topKFrequent_postcond (nums : List Int) (k : Nat) (result: List Int) (h_precond : topKFrequent_precond (nums) (k)) : Prop :=
  -- !benchmark @start postcond
  -- Result contains exactly k elements
  result.length = k ∧

  -- All elements in result appear in the original list
  result.all (· ∈ nums) ∧

  -- All elements in result are distinct
  List.Pairwise (· ≠ ·) result ∧

  -- For any element in result and any element not in result, the frequency of the
  -- element in result is greater or equal (simplified: no tie-breaking constraint on ordering)
  (result.all (fun x =>
    nums.all (fun y =>
      y ∈ result ∨ nums.count x ≥ nums.count y
    ))) ∧

  -- Elements in result are ordered by non-increasing frequency
  List.Pairwise (fun (x, i) (y, j) =>
    i < j → nums.count x ≥ nums.count y
  ) result.zipIdx
  -- !benchmark @end postcond


-- !benchmark @start proof_aux

-- !benchmark @end proof_aux


theorem topKFrequent_spec_satisfied (nums: List Int) (k: Nat) (h_precond : topKFrequent_precond (nums) (k)) :
    topKFrequent_postcond (nums) (k) (topKFrequent (nums) (k) h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof




Tests:
input: nums: [1, 1, 1, 2, 2, 3]
k: 2
nums: [4, 1, -1, 2, -1, 2, 3]
k: 2
nums: [5]
k: 1
nums: [7, 7, 7, 8, 8, 9]
k: 1
k: 0
expected: [1, 2]
[-1, 2]
[5]
[7]
unexpected: [1, 3]
[2, 3]
[-1, 4]
[4, 3]
[8]
[0]

Inputs rejected by the precondition:
nums: [1, 2, 3]
k: 4