Stable marriage problem algorithm implementation in Elixir

stable_marriage_problem.ex

defmodule StableMarriageProblem do
  @moduledoc """
  Implements the algorithm for stable marriage problem.
  """

  @type person_id :: integer()

  # Denotes a preference list for a single person
  @type preference_list :: [person_id(), ...]

  # Denotes preference lists of multiple people where the key is the id of the person
  @type preference_map :: %{required(person_id()) => preference_list()}

  @type proposals :: %{proposer_id :: integer() => proposed_id :: integer()}

  # Denotes matches between men & women
  @type matches :: %{person_id() => person_id()}

  # Change this number based on how many people you want to match
  @num_men 4

  # Preference map for men
  @men (1..@num_men
    |> Enum.to_list()
    |> Map.new(& {&1, Enum.shuffle(1..@num_men)}))

  # Preference map for women
  @women (1..@num_men
    |> Enum.to_list()
    |> Map.new(& {&1, Enum.shuffle(1..@num_men)}))

  IO.puts("== MEN PREFERENCE MAP ==")
  Enum.map(@men, & IO.inspect/1)

  IO.puts("\n== WOMEN PREFERENCE MAP ==")
  Enum.map(@women, & IO.inspect/1)

  def stabilise do
    men_preference_map = @men

    # When we're starting nobody is matched
    non_matched_men = Map.keys(men_preference_map)

    # Initialize matches map
    matches = Map.new(non_matched_men, & {&1, nil})

    do_stabilise(matches, men_preference_map, non_matched_men)
  end

  def do_stabilise(matches, _, _non_matched_men = []) do
    IO.puts("\n== Matching complete ==")
    Enum.map(matches, & IO.inspect/1)
  end

  def do_stabilise(matches, men_preference_list, non_matched_men) do
    {updated_preference_list, proposals} = get_proposals(men_preference_list, non_matched_men)

    IO.puts("\n== Proposals ==")
    Enum.map(proposals, &IO.inspect/1)
    IO.puts("")

    new_matches = evaluate_proposals(proposals, matches)

    IO.puts("\n== Matches ==")
    Enum.map(new_matches, &IO.inspect/1)
    IO.puts("")

    # Find remaining non-matched men
    non_matched_men =
      new_matches
      |> find_non_matched_men()
      |> IO.inspect(label: "\nNON MATCHED MEN")

    do_stabilise(new_matches, updated_preference_list, non_matched_men)
  end

  def find_non_matched_men(matches) do
    matches
    |> Map.reject(fn {_k, v} -> not is_nil(v) end)
    |> Map.keys()
  end

  # Returns proposals for each proposing men, and updated preference map
  # We basically drop their first preference from their preference list and update preference map
  @spec get_proposals(preference_map(), proposing_men :: [integer()]) :: {preference_map(), proposals()}
  def get_proposals(men_preference_map, proposing_men) do
    Enum.reduce(proposing_men, {men_preference_map, _proposals = %{}},
      fn man_id, {men_preference_map, proposals}  ->
        {propose_to, preference_list} =
          men_preference_map
          |> Map.get(man_id)
          |> List.pop_at(0)

        proposals = Map.put(proposals, man_id, propose_to)
        men_preference_map = Map.put(men_preference_map, man_id, preference_list)

        {men_preference_map, proposals}
    end)
  end

  # Evaluates proposals and returns new matches
  def evaluate_proposals(proposals, matches) do
    Enum.reduce(proposals, matches, fn
      {proposing_man_id, woman_id}, matches ->
        case fetch_match_for_woman(matches, woman_id) do
          # When woman isn't matched with anyone yet
          nil ->
            IO.puts("W##{woman_id} accepts the proposal of M##{proposing_man_id}")
            Map.put(matches, proposing_man_id, woman_id)

          # When woman already has a match
          {matched_man_id, _woman_id} ->
            if woman_accepts_proposal?(woman_id, matched_man_id, proposing_man_id) do
              IO.puts("W##{woman_id} accepts the proposal of M##{proposing_man_id}")
              matches
              |> Map.put(matched_man_id, nil)
              |> Map.put(proposing_man_id, woman_id)

            else
              IO.puts("W##{woman_id} does NOT accept the proposal of M##{proposing_man_id}")
              matches
            end
        end
    end)
  end

  # Returns the match for a given woman, if there's no match returns nil
  @spec fetch_match_for_woman(matches(), person_id()) :: {person_id(), person_id()} | nil
  def fetch_match_for_woman(matches, woman_id) do
    Enum.find(matches, fn
      {_man_id, ^woman_id} -> true
      _ -> false
    end)
  end

  # A woman accepts a proposal only if the proposing man is higher on woman's preference list
  # than the man she is currently matched.
  @spec woman_accepts_proposal?(person_id(), person_id(), person_id()) :: boolean()
  def woman_accepts_proposal?(woman_id, matched_man_id, proposing_man_id) do
    @women
    |> Map.get(woman_id)
    |> Enum.find(& &1 == matched_man_id || &1 == proposing_man_id)
    |> then(& &1 == proposing_man_id)
  end
end

StableMarriageProblem.stabilise()

I could probably just keep a list of non_matched_men and push into it instead of going through the matches again and again to find non-matched men.