Minimal sub arborescence

I was playing with this during lunch and came up with this:

defmodule Arby do
  @doc """
  Finds the minimal sub arborescence containing all the vertices in `vertices` of the given arborescence `graph`.
  """
  @spec minimal_sub_arborscence(Graph.t(), Graph.label()) :: Graph.t()
  def minimal_sub_arborscence(graph, vertices) do
    if Graph.is_arborescence?(graph) do
      do_minimal_sub_arborscence(graph, vertices)
    else
      raise("Graph should be arborescence!")
    end
  end

  defp do_minimal_sub_arborscence(graph, vertices) do
    vertices
    |> Enum.map(&Graph.Utils.vertex_id/1)
    |> then(&fully_connected_vertex_ids(graph, &1))
    |> MapSet.to_list()
    |> then(&Map.take(graph.vertices, &1))
    |> Map.values()
    |> then(&Graph.subgraph(graph, &1))
  end

  @doc """
  Considers the given arborescence `graph` and the given list of vertex_ids,
  returning a MapSet of vertex_ids representing the minimal set
  required for all vertexes to be connected.
  """
  @spec fully_connected_vertex_ids(Graph.t(), [Graph.vertex_id()]) :: MapSet.t()
  def fully_connected_vertex_ids(graph, vertex_ids) do
    fully_connected_vertex_ids(graph, vertex_ids, MapSet.new(vertex_ids))
  end

  @spec fully_connected_vertex_ids(Graph.t(), [Graph.vertex_id()], MapSet.t()) :: MapSet.t()
  def fully_connected_vertex_ids(_graph, [], accumulated_vertex_ids) do
    accumulated_vertex_ids
  end

  def fully_connected_vertex_ids(
        graph,
        [vertex_id | remaining_vertex_ids],
        accumulated_vertex_ids
      ) do
    fully_connected_vertex_ids(
      graph,
      remaining_vertex_ids,
      MapSet.union(accumulated_vertex_ids, Map.get(graph.in_edges, vertex_id))
    )
  end
end

I think it’s better because I am fully exploiting the arborescence structure. In such a structure, every “vertex” has exactly 1 parent. So the fully_connected_vertex_ids/2 function just uses graph.in-edges to “get the parent” of each vertex we care about, until we have a final MapSet of vertex_ids representing a fully connected set of vertexes.

Once you have this, it’s a simple matter of making a subgraph.

Here’s my test module for that code:

defmodule ArbyTest do
  use ExUnit.Case
  doctest Arby

  @doc """
  It looks like this:

      a
    / \
    b   c
  /|\   \
  d e f   g

  """
  def example_graph() do
    Graph.new()
    |> Graph.add_edge(:a, :b)
    |> Graph.add_edge(:b, :d)
    |> Graph.add_edge(:b, :e)
    |> Graph.add_edge(:b, :f)
    |> Graph.add_edge(:a, :c)
    |> Graph.add_edge(:c, :g)
  end

  describe "minimal_sub_arborscence/2" do
    test "gets minimal_graph with reduced set of vetices" do
      graph = example_graph()

      for {input, expected_output} <- [
            {[:c, :d, :e], [:a, :b, :c, :d, :e]},
            {[:d, :f], [:b, :d, :f]}
          ] do
        assert graph
               |> Arby.minimal_sub_arborscence(input)
               |> Map.get(:vertices)
               |> Map.values() == expected_output
      end
    end
  end

  describe "arborescence assumptions" do
    test "arborescent graphs permit only 1 parent per vertex" do
      graph =
        example_graph()
        |> Graph.add_edge(:z, :b)

      refute Graph.is_arborescence?(graph)
    end
  end
end