Extend Map parser

2023/d11.py

@@ -29,7 +29,7 @@ class Day11(aoc.Challenge):
 
     def solver(self, puzzle_input: InputType, part_one: bool) -> int:
         """Return the sum of the distances between shifted galaxies."""
-        data = puzzle_input.coords["#"]
+        data = puzzle_input["#"]
         distance = 1 if part_one else 1_000_000 - 1
         min_x, min_y, max_x, max_y = aoc.bounding_coords(data)
 

2024/d04.py

@@ -28,10 +28,10 @@ def part1(self, puzzle_input: aoc.Map) -> int:
         """Count occurances of XMAS in the word search."""
         return sum(
             all(
-                (start + distance * direction) in puzzle_input.coords[letter]
+                (start + distance * direction) in puzzle_input[letter]
                 for distance, letter in enumerate("MAS", start=1)
             )
-            for start in puzzle_input.coords["X"]
+            for start in puzzle_input["X"]
             for direction in aoc.EIGHT_DIRECTIONS
         )
 
@@ -45,10 +45,10 @@ def part2(self, puzzle_input: aoc.Map) -> int:
         }
         return sum(
             all(
-                (start + offset * 1j ** rotation) in puzzle_input.coords[letter]
+                (start + offset * 1j ** rotation) in puzzle_input[letter]
                 for offset, letter in want.items()
             )
-            for start in puzzle_input.coords["A"]
+            for start in puzzle_input["A"]
             for rotation in range(4)
         )
 

2024/d06.py

@@ -43,9 +43,9 @@ def walk(
 
     def solver(self, puzzle_input: aoc.Map, part_one: bool) -> int:
         all_spots = puzzle_input.all_coords
-        blocked = puzzle_input.coords["#"]
+        blocked = puzzle_input["#"]
         start_pos, start_dir = next(
-            (puzzle_input.coords[arrow].copy().pop(), aoc.ARROW_DIRECTIONS[arrow])
+            (puzzle_input[arrow].copy().pop(), aoc.ARROW_DIRECTIONS[arrow])
             for arrow in "<>v^"
             if arrow in puzzle_input.coords
         )

2024/d08.py

@@ -33,7 +33,7 @@ def solver(self, puzzle_input: aoc.Map, part_one: bool) -> int:
         all_locations = puzzle_input.all_coords
         antinodes = set()
         for freq in puzzle_input.non_blank_chars:
-            for a, b in itertools.combinations(puzzle_input.coords[freq], 2):
+            for a, b in itertools.combinations(puzzle_input[freq], 2):
                 sequences: tuple[collections.abc.Iterable[int], collections.abc.Iterable[int]]
                 if part_one:
                     sequences = ([2], [-1])

pylib/parsers.py

@@ -295,19 +295,26 @@ def parse(self, puzzle_input: str) -> set[complex]:
 class Map:
     max_x: int
     max_y: int
-    chars: dict[complex, str]
-    coords: dict[str, set[complex]]
+    chars: dict[complex, str | int]
+    coords: dict[str | int, set[complex]]
     all_coords: set[complex]
     blank_char: str
     non_blank_chars: set[str]
 
-    @property
-    def width(self) -> int:
-        return self.max_x + 1
-
-    @property
-    def height(self) -> int:
-        return self.max_y + 1
+    def __post_init__(self) -> None:
+        self.width = self.max_x + 1
+        self.height = self.max_y + 1
+
+    def __getitem__(self, key: str | complex) -> str | int | set[complex] | list[set[complex]]:
+        if isinstance(key, complex):
+            return self.chars[key]
+        if isinstance(key, int):
+            return self.coords[key]
+        if isinstance(key, str) and len(key) == 1:
+            return self.coords[key]
+        if isinstance(key, collections.abc.Sequence):
+            return [self.coords[k] for k in key]
+        raise ValueError(f"Could not index on {key}")
 
     @property
     def size(self) -> int: