Breadth First Search: Shortest Reach

shortest_reach.ex

defmodule ShortestReach do
  def main do
    IO.gets("")
    |> String.trim()
    |> String.to_integer()
    |> shortest_reach()
  end

  defp shortest_reach(0) do
  end

  defp shortest_reach(q) do
    [n, m] =
      IO.gets("")
      |> String.trim()
      |> String.split(" ")
      |> Enum.map(&String.to_integer/1)

    edges = read_edges(m)

    s =
      IO.gets("")
      |> String.trim()
      |> String.to_integer()

    init_graph(n)
    |> add_edges(edges)
    |> bfs(s)
    |> expand_paths(n, s)
    |> Enum.reverse()
    |> Enum.join(" ")
    |> IO.puts()

    shortest_reach(q - 1)
  end

  defp read_edges(0), do: []

  defp read_edges(m) do
    [
      IO.gets("")
      |> String.trim()
      |> String.split(" ")
      |> Enum.map(&String.to_integer/1)
      | read_edges(m - 1)
    ]
  end

  defp init_graph(0), do: %{0 => MapSet.new()}
  defp init_graph(n), do: init_graph(n - 1) |> Map.put(n, MapSet.new())

  defp add_edges(graph, []), do: graph

  defp add_edges(graph, [[u, v] | edges_tail]) do
    Map.put(graph, u, MapSet.put(graph[u], v))
    |> Map.put(v, MapSet.put(graph[v], u))
    |> add_edges(edges_tail)
  end

  defp bfs(graph, s), do: bfs(%{s: 0}, graph, MapSet.to_list(graph[s]))

  defp bfs(paths, graph, neighbors, next_nodes \\ [], layer \\ 1)

  defp bfs(paths, _graph, [], [], _layer), do: paths

  defp bfs(paths, graph, [], next_nodes, layer),
    do: bfs(paths, graph, next_nodes, [], layer + 1)

  defp bfs(paths, graph, [u | neighbors], next_nodes, layer)
       when is_map_key(paths, u),
       do: bfs(paths, graph, neighbors, next_nodes, layer)

  defp bfs(paths, graph, [u | neighbors], next_nodes, layer),
    do:
      Map.put_new(paths, u, layer)
      |> bfs(graph, neighbors, MapSet.to_list(graph[u]) ++ next_nodes, layer)

  defp expand_paths(_paths, 0, _s), do: []

  defp expand_paths(paths, s, s), do: expand_paths(paths, s - 1, s)

  defp expand_paths(paths, n, s) when is_map_key(paths, n),
    do: [paths[n] * 6] ++ expand_paths(paths, n - 1, s)

  defp expand_paths(paths, n, s),
    do: [-1] ++ expand_paths(paths, n - 1, s)
end

ShortestReach.main()