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(FYS) :: {edge, integer(), integer(), FYS}.

-file("src/yog/mst.gleam", 52).
-spec do_kruskal(
    list(edge(FYZ)),
    yog@disjoint_set:disjoint_set(integer()),
    list(edge(FYZ))
) -> list(edge(FYZ)).
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(), FYU),
    fun((FYU, FYU) -> gleam@order:order())
) -> list(edge(FYU)).
kruskal(Graph, Compare) ->
    Node_ids = maps:keys(erlang:element(3, Graph)),
    Edges = begin
        _pipe = maps:to_list(erlang:element(4, Graph)),
        _pipe@2 = gleam@list:flat_map(
            _pipe,
            fun(Entry) ->
                {From_id, Targets} = Entry,
                _pipe@1 = maps:to_list(Targets),
                gleam@list:filter_map(
                    _pipe@1,
                    fun(Target) ->
                        {To_id, Weight} = Target,
                        case (erlang:element(2, Graph) =:= undirected) andalso (From_id
                        > To_id) of
                            true ->
                                {error, nil};

                            false ->
                                {ok, {edge, From_id, To_id, Weight}}
                        end
                    end
                )
            end
        ),
        gleam@list:sort(
            _pipe@2,
            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, []).