aoc2020day7.ex

defmodule Rule do
  defstruct [:from, :to, :count]
end

defmodule Day7 do
  def read_input do
    File.read!("input/day7.txt")
    |> String.split("\n", trim: true)
    |> Enum.map(&(String.replace(&1, "bags", "bag")))
    |> Enum.map(&(String.trim(&1, ".")))
    |> Enum.flat_map(&parse_rule/1)
  end

  def search_with_counts(bags, start \\ "shiny gold bag", counted \\ 1) do
    if Map.has_key?(bags, start) do
      inner_bags = Map.get(bags, start)
                   |> Enum.map(fn rule -> search_with_counts(bags, rule.to, counted * rule.count) end)
      
      Enum.sum(inner_bags) + counted
    else
      counted
    end
  end

  def search(bags, start \\ "shiny gold bag") do
    Map.get(bags, start, [])
    |> Enum.map(&(search(bags, &1)))
    |> Enum.reduce(MapSet.new([start]), &MapSet.union/2 )
  end

  def parse_rule(str) do
    [src, destinations] = String.split(str, " contain ")
    destinations = String.split(destinations, ", ")
                   |> Enum.map(&Integer.parse/1)

    for { count, " " <> bag_type } <- destinations do
      %Rule{to: bag_type, from: src, count: count}
    end
  end

  def part1 do
    res = read_input()
          |> Enum.reduce(%{}, fn rule, acc ->
            Map.update(
              acc,
              rule.to,
              MapSet.new([rule.from]),
              &(MapSet.put(&1, rule.from)))
          end)
          |> search()
          |> Enum.count
    IO.inspect(res - 1)
  end

  def part2 do
    res = read_input()
          |> Enum.reduce(%{}, fn rule, acc ->
            Map.update(
              acc,
              rule.from,
              MapSet.new([rule]),
              &(MapSet.put(&1, rule)))
          end)
          |> search_with_counts

    IO.inspect(res - 1)
  end
end