day 02: solve part 1

day_02/program.py

@@ -2,14 +2,57 @@
 
 # https://adventofcode.com/2023/day/2
 
+import re
+
+
 def get_lines(filename: str) -> list:
     with open(filename, "r") as file:
         return [line.strip() for line in file.readlines()]
 
 
+def get_games(lines: list) -> dict:
+    games = {}
+    for line in lines:
+        game_number, grabs = line.split(": ")
+        game_number = game_number[5:]  # shave off "Game " from string
+        games[int(game_number)] = grabs
+    return games
+
+
+def get_max_cubes(grabs: str) -> dict:
+    max_cubes = {}
+    for grab in grabs.split(";"):
+        for cube_string in grab.split(", "):
+            num_cubes_str, colour = cube_string.split()
+            num_cubes = int(num_cubes_str)
+            if colour not in max_cubes or num_cubes > max_cubes[colour]:
+                max_cubes[colour] = num_cubes
+    return max_cubes
+
+
+def game_is_valid(game: dict, cube_pools: dict) -> bool:
+    for colour in game.keys():
+        if colour not in cube_pools or game[colour] > cube_pools[colour]:
+            return False
+    return True
+
+
+def sum_of_valid_games(games: dict, cube_pools: dict) -> int:
+    valid_games = 0
+    for game_num in games.keys():
+        grabs = games[game_num]
+        max_cubes = get_max_cubes(grabs)
+        if game_is_valid(max_cubes, cube_pools):
+            valid_games += game_num
+    return valid_games
+
+
 def main():
-    lines = get_lines("test-input.txt")
-    # lines = get_lines("input.txt")
+    lines = get_lines("input.txt")
+    cube_pools = {"red": 12, "green": 13, "blue": 14}
+    games = get_games(lines)
+    result = sum_of_valid_games(games, cube_pools)
+    print(f"Part 1: The sum of valid games is: {result}")
 
 
 if __name__ == '__main__':

day_02/test_program.py

@@ -8,6 +8,43 @@ class TestThing(unittest.TestCase):
     def setUp(self):
         pass
 
+    def testGetGames(self):
+        lines = program.get_lines("test-input.txt")
+        games = program.get_games(lines)
+        target = {
+            1: "3 blue, 4 red; 1 red, 2 green, 6 blue; 2 green",
+            2: "1 blue, 2 green; 3 green, 4 blue, 1 red; 1 green, 1 blue",
+            3: "8 green, 6 blue, 20 red; 5 blue, 4 red, 13 green; 5 green, 1 red",
+            4: "1 green, 3 red, 6 blue; 3 green, 6 red; 3 green, 15 blue, 14 red",
+            5: "6 red, 1 blue, 3 green; 2 blue, 1 red, 2 green"
+        }
+        self.assertEqual(games, target)
+
+    def testMaxCubes(self):
+        test_cases = [
+            ("3 blue, 2 red; 3 red, 4 blue", {"blue": 4, "red": 3}),
+            ("2 red, 1 blue; 4 blue, 2 yellow", {"red": 2, "blue": 4, "yellow": 2})
+        ]
+        for grabs, max_cubes in test_cases:
+            self.assertEqual(program.get_max_cubes(grabs), max_cubes)
+
+    def testGameIsValid(self):
+        cube_pools = {"red": 4, "green": 5, "blue": 6}
+        test_cases = [
+            ({"blue": 4, "red": 3, "green": 1}, True),
+            ({"blue": 7}, False),
+            ({"green": 2, "red": 1}, True),
+            ({"blue": 3, "red": 2, "yellow": 1}, False)
+        ]
+        for test_case, test_result in test_cases:
+            self.assertEqual(program.game_is_valid(test_case, cube_pools), test_result)
+
+    def testSumOfValidGames(self):
+        lines = program.get_lines("test-input.txt")
+        games = program.get_games(lines)
+        cube_pools = {"red": 12, "green": 13, "blue": 14}
+        self.assertEqual(program.sum_of_valid_games(games, cube_pools), 8)
+
 
 if __name__ == '__main__':
     unittest.main()