Solve 2021 day 6

farbfetzen · farbfetzen · commit cc327f3d1871 · 2021-12-06T20:04:47.000+01:00
diff --git a/README.md b/README.md
@@ -3,7 +3,7 @@
 
 | Year | Stars |
 |------|------:|
-| 2021 |  10 ⭐ |
+| 2021 |  12 ⭐ |
 | 2020 |  50 ⭐ |
 | 2019 |  44 ⭐ |
 | 2018 |   8 ⭐ |
diff --git a/input/2021-06-sample.txt b/input/2021-06-sample.txt
@@ -0,0 +1 @@
+3,4,3,1,2
diff --git a/python/2021/day06.py b/python/2021/day06.py
@@ -0,0 +1,32 @@
+# https://adventofcode.com/2021/day/6
+
+
+def get_data(filename):
+    with open(filename) as file:
+        return [int(x) for x in file.read().split(",")]
+
+
+def simulate(fish_ages):
+    count = [fish_ages.count(i) for i in range(9)]
+    result_at_day_80 = 0
+    for day in range(1, 256+1):
+        zeros = count[0]
+        count[:-1] = count[1:]
+        count[6] += zeros
+        count[8] = zeros
+        if day == 80:
+            result_at_day_80 = sum(count)
+    return result_at_day_80, sum(count)
+
+
+sample_data = get_data("../../input/2021-06-sample.txt")
+challenge_data = get_data("../../input/2021-06-input.txt")
+
+if __name__ == "__main__":
+    sample_part_1, sample_part_2 = simulate(sample_data)
+    assert sample_part_1 == 5934
+    assert sample_part_2 == 26984457539
+
+    challenge_part_1, challenge_part_2 = simulate(challenge_data)
+    print(challenge_part_1)  # 372300
+    print(challenge_part_2)  # 1675781200288
diff --git a/python/2021/test2021.py b/python/2021/test2021.py
@@ -5,6 +5,7 @@
 import day03
 import day04
 import day05
+import day06
 
 
 class Test2021(unittest.TestCase):
@@ -47,6 +48,15 @@ def test_05(self):
 
         self.assertEqual(day05.part_1(day05.challenge_data), 6267)
         self.assertEqual(day05.part_2(day05.challenge_data), 20196)
+        
+    def test_06(self):
+        sample_part_1, sample_part_2 = day06.simulate(day06.sample_data)
+        self.assertEqual(sample_part_1, 5934)
+        self.assertEqual(sample_part_2, 26984457539)
+
+        challenge_part_1, challenge_part_2 = day06.simulate(day06.challenge_data)
+        self.assertEqual(challenge_part_1, 372300)
+        self.assertEqual(challenge_part_2, 1675781200288)
 
 
 if __name__ == "__main__":
