Finish part 2 of day 7

1 files changed, 36 insertions, 6 deletions

diff --git a/aoc-2020-gleam/src/days/day07.gleam b/aoc-2020-gleam/src/days/day07.gleam
index eee817f..c7c26dc 100644
--- a/aoc-2020-gleam/src/days/day07.gleam
+++ b/aoc-2020-gleam/src/days/day07.gleam
@@ -2,7 +2,8 @@ import gleam/io
 import gleam/list
 import gleam/function as fun
 import gleam/pair
-import gleam/iterator as iter
+import gleam/result as res
+import gleam/iterator.{Iterator} as iter
 import gleam/map.{Map}
 import ext/resultx as resx
 import ext/iteratorx as iterx
@@ -12,11 +13,17 @@ import util/parser as p
 
 const special_bag = "shiny gold"
 
+type BagId =
+  String
+
 type BagEdge =
-  #(String, Int)
+  #(BagId, Int)
 
 type BagGraph =
-  Map(String, List(BagEdge))
+  Map(BagId, List(BagEdge))
+
+type BagNeighbourFun =
+  fn(BagId) -> Iterator(BagId)
 
 fn parse_graph(lines: List(String)) -> BagGraph {
   let bag_type_parser =
@@ -48,15 +55,32 @@ fn parse_graph(lines: List(String)) -> BagGraph {
   |> map.from_list
 }
 
-fn part1(lines: List(String)) -> Int {
-  let graph = parse_graph(lines)
-  let neighbours = fn(bag) {
+fn neighbour_fun(graph: BagGraph) -> BagNeighbourFun {
+  fn(bag) {
     graph
     |> map.get(bag)
     |> resx.force_unwrap
     |> list.map(with: pair.first)
     |> iter.from_list
   }
+}
+
+fn bag_count(of bag: BagId, in graph: BagGraph) -> Int {
+  list.fold(
+    over: graph
+    |> map.get(bag)
+    |> res.unwrap(or: []),
+    from: 1,
+    with: fn(sum, edge) {
+      let #(next_bag, next_count) = edge
+      sum + next_count * bag_count(of: next_bag, in: graph)
+    },
+  )
+}
+
+fn part1(lines: List(String)) -> Int {
+  let graph = parse_graph(lines)
+  let neighbours = neighbour_fun(graph)
 
   graph
   |> map.keys
@@ -69,12 +93,18 @@ fn part1(lines: List(String)) -> Int {
   })
 }
 
+fn part2(lines: List(String)) -> Int {
+  bag_count(of: special_bag, in: parse_graph(lines)) - 1
+}
+
 pub fn run() -> Nil {
   let test = input_util.read_lines("test07")
   assert 4 = part1(test)
+  assert 32 = part2(test)
 
   let input = input_util.read_lines("day07")
   io.debug(part1(input))
+  io.debug(part2(input))
 
   Nil
 }