yog

-module(yog@mst).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/yog/mst.gleam").
-export([kruskal/2]).
-export_type([edge/1]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

-type edge(GXQ) :: {edge, integer(), integer(), GXQ}.

-file("src/yog/mst.gleam", 53).
-spec do_kruskal(
    list(edge(GXX)),
    yog@disjoint_set:disjoint_set(integer()),
    list(edge(GXX))
) -> list(edge(GXX)).
do_kruskal(Edges, Disjoint_set_state, Acc) ->
    case Edges of
        [] ->
            lists:reverse(Acc);

        [Edge | Rest] ->
            {Disjoint_set1, Root_from} = yog@disjoint_set:find(
                Disjoint_set_state,
                erlang:element(2, Edge)
            ),
            {Disjoint_set2, Root_to} = yog@disjoint_set:find(
                Disjoint_set1,
                erlang:element(3, Edge)
            ),
            case Root_from =:= Root_to of
                true ->
                    do_kruskal(Rest, Disjoint_set2, Acc);

                false ->
                    Next_disjoint_set = yog@disjoint_set:union(
                        Disjoint_set2,
                        erlang:element(2, Edge),
                        erlang:element(3, Edge)
                    ),
                    do_kruskal(Rest, Next_disjoint_set, [Edge | Acc])
            end
    end.

-file("src/yog/mst.gleam", 25).
?DOC(
    " Finds the Minimum Spanning Tree (MST) using Kruskal's algorithm.\n"
    "\n"
    " Returns a list of edges that form the MST. The total weight of these edges\n"
    " is minimized while ensuring all nodes are connected.\n"
    "\n"
    " **Time Complexity:** O(E log E) where E is the number of edges\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " let mst_edges = mst.kruskal(in: graph, with_compare: int.compare)\n"
    " // => [Edge(1, 2, 5), Edge(2, 3, 3), ...]\n"
    " ```\n"
).
-spec kruskal(
    yog@model:graph(any(), GXS),
    fun((GXS, GXS) -> gleam@order:order())
) -> list(edge(GXS)).
kruskal(Graph, Compare) ->
    Node_ids = maps:keys(erlang:element(3, Graph)),
    Edges = begin
        _pipe = gleam@dict:fold(
            erlang:element(4, Graph),
            [],
            fun(Acc, From_id, Targets) ->
                Inner_edges = gleam@dict:fold(
                    Targets,
                    [],
                    fun(Inner_acc, To_id, Weight) ->
                        case (erlang:element(2, Graph) =:= undirected) andalso (From_id
                        > To_id) of
                            true ->
                                Inner_acc;

                            false ->
                                [{edge, From_id, To_id, Weight} | Inner_acc]
                        end
                    end
                ),
                lists:append([Inner_edges, Acc])
            end
        ),
        gleam@list:sort(
            _pipe,
            fun(A, B) -> Compare(erlang:element(4, A), erlang:element(4, B)) end
        )
    end,
    Initial_disjoint_set = gleam@list:fold(
        Node_ids,
        yog@disjoint_set:new(),
        fun yog@disjoint_set:add/2
    ),
    do_kruskal(Edges, Initial_disjoint_set, []).