stable matching in Lean

StableMatching.lean

-- Extended Stable Matching for Multiple Positions and Quotas
import Cpnu.Data

namespace ExtendedGaleShapley

structure Candidate where
  id          : Nat
  preferences : List Nat
deriving Repr, Inhabited


structure Post where
  id          : Nat
  capacity    : Nat
  preferences : List Nat
deriving Repr, Inhabited


structure MatchState where
  free        : List Nat
  proposals   : Std.HashMap (Nat × Nat) Bool
  assignments : Std.HashMap Nat (List Nat)
deriving Repr, Inhabited


abbrev GS := StateM MatchState


def markProposed (cid pid : Nat) : GS Unit  := do
  let s ← get
  let newProps := s.proposals.insert (cid, pid) true
  set { s with proposals := newProps }


def alreadyProposed (cid pid : Nat) : GS Bool := do
  let s ← get
  pure <| s.proposals.getD (cid, pid) false


def insertByPreference
  (accepted : List Nat) (newCid : Nat) (post : Post)
  : List Nat × List Nat :=
  let all := accepted ++ [newCid]
  let ordered := all.mergeSort (fun a b =>
    let aIdx := post.preferences.idxOf? a
    let bIdx := post.preferences.idxOf? b
    match aIdx, bIdx with
    | some n, some m => n < m
    | _, _ => false)
  ordered.splitAt post.capacity


partial def galeShapley
  (candidates : Array Candidate)
  (posts : Array Post) : GS Unit := do
  let mut done := false
  while !done do
    let s ← get
    match s.free with
    | [] =>
      -- no more free candidate
      done := true
    | cid :: rest =>
      dbg_trace "candidate {cid} selected from {s.free.length} free"
      -- Find next post to try
      let cand := candidates.find? (fun c => c.id == cid) |>.get!
      let mut targetPid : Option Nat := none

      -- Find first post not yet proposed to
      for pid in cand.preferences do
        let already ← alreadyProposed cid pid
        if !already then
          targetPid := some pid
          break

      match targetPid with
      | none =>
          -- No more posts to try for this candidate
          set { s with free := rest }
      | some pid =>
          -- Mark as proposed
          markProposed cid pid

          -- Get current assignments
          let post := posts.find? (fun p => p.id == pid) |>.get!
          let s' ← get  -- Get updated state after marking proposal
          let accepted := s'.assignments.getD pid []

          -- Check if there's space
          if accepted.length < post.capacity then
            -- Accept directly
            let (newAccepted, _) := insertByPreference accepted cid post
            let newAssignments := s'.assignments.insert pid newAccepted
            set { s' with
              free := rest,
              assignments := newAssignments }
          else
            -- Select best subset from (accepted + new candidate)
            let (newAccepted, rejected) :=
              insertByPreference accepted cid post
            let newAssignments := s'.assignments.insert pid newAccepted
            set { s' with
              free := rest ++ rejected,
              assignments := newAssignments }


-- tests

def candidates : Array Candidate := #[
  { id := 0, preferences := [0, 1] },
  { id := 1, preferences := [0, 1] },
  { id := 2, preferences := [0, 1] }
]

def posts : Array Post := #[
  { id := 0, capacity := 2, preferences := [2, 1, 0] },
  { id := 1, capacity := 1, preferences := [0, 1, 2] }
]

-- load data

def buildInstance (dataDir : String) (b : Nat)
  : IO (Array Candidate × Array Post) := do

  let (candidates, jobs, ranking) ← Data.loadData dataDir b

  let candidates : Array Candidate :=
    candidates.toArray.map (fun c =>
      { id := c.id , preferences := c.preferences} )

  let posts : Array Post := jobs.toArray.map
    (fun j =>
      let cs := candidates.toList.filter
        (fun c => c.preferences.contains j.code)
      let rs := cs.map (fun c => (c.id, ranking[(j.code, c.id)]!))
        |>.mergeSort (fun a b => a.2 ≥ b.2) |>.map (·.1)
      { id := j.code,
        capacity  := j.total,
        preferences := rs })

  return (candidates, posts)

def initialState (cs : Array Candidate) : MatchState := {
  free := cs.toList.map (·.id),
  proposals := Std.HashMap.emptyWithCapacity 100,
  assignments := Std.HashMap.emptyWithCapacity 100
}


end ExtendedGaleShapley