Adding day06 part 2

puzzles/day06.py

@@ -8,47 +8,45 @@ Run with:
 import pathlib
 import sys
 
-from dataclasses import dataclass
+from collections import Counter
 from typing import List, Tuple
 
 
-@dataclass
-class LanternFish:
-    reproduce_timer: int
-    children: int = 0
-
-class FishSchool:
-    fish = None
+class Reproduce:
+    times = None
 
     def __init__(self, times: List[int]) -> None:
-        self.fish = []
-        for time in times:
-            self.fish.append(LanternFish(reproduce_timer=time))
+        self.times = times
 
-    def simulate(self, days: int=80) -> None:
+    def start(self, days: int=80) -> None:
+        counts = Counter(self.times)
         for day in range(days):
-            self._check_school()
+            next_counts = Counter() # Can't modify while iterating
+
+            for count_key, count in counts.items():
+                if count_key == 0:
+                    next_counts[6] += count # Move the parents back to 6 days
+                    next_counts[8] += count # Add the new fish
+                else:
+                    next_counts[count_key - 1] += count # Move count to next day count
+
+            counts = next_counts
+        self.time_counts = counts
 
-    def _check_school(self) -> None:
-        for index, fish in enumerate(self.fish):
-            if fish.reproduce_timer == 0:
-                self.fish.append(LanternFish(reproduce_timer=9))    # Need to correct (should be 8)
-                self.fish[index].reproduce_timer = 6
-                self.fish[index].children += 1
-            else:
-                self.fish[index].reproduce_timer -= 1
+    def count_fish(self):
+        return sum(self.time_counts.values())
 
 
 def part1(inputs: List[int]) -> int:
-    school = FishSchool(inputs)
-    school.simulate(days=80)
-    return len(school.fish)
+    school = Reproduce(inputs)
+    school.start(days=80)
+    return school.count_fish()
 
 
 def part2(inputs: List[int]) -> int:
-    school = FishSchool(inputs)
-    school.simulate(days=256)
-    return len(school.fish)
+    school = Reproduce(inputs)
+    school.start(days=256)
+    return school.count_fish()
 
 
 def parse(inputs: str) -> List[int]:

puzzles/test_day06.py

@@ -24,11 +24,9 @@ def test_example2(example_data):
     assert aoc.part2(example_data) == 26984457539
 
 
-@pytest.mark.skip(reason="Not implemented")
 def test_part1(day06_data):
-    assert aoc.part1(day06_data) == ...
+    assert aoc.part1(day06_data) == 386755
 
 
-@pytest.mark.skip(reason="Not implemented")
 def test_part2(day06_data):
-    assert aoc.part2(day06_data) == ...
+    assert aoc.part2(day06_data) == 1732731810807