AoC 2024 day 4 part 2

aoc2024/src/day04/python/solution.py

@@ -14,6 +14,10 @@
 from typing import Sequence
 
 def _extract_all_lines(input: Sequence[str]) -> Sequence[str]:
+    """Breaks down the input grid into lines to facilitate word search.
+    Time complexity: O(n) + O(nm) + O(2m*n) + O(2m*n) = O(nm)
+    Space complexity: O(n) + O(m) + O(m+n) + O(m+n) = O(3m+3n) = O(m+n)
+    """
     lines = []
     width = len(input[0])
     height = len(input)
@@ -23,7 +27,7 @@ def _extract_all_lines(input: Sequence[str]) -> Sequence[str]:
     # Verticals
     for x in range(width):
         vertical = []
-        for y in range(len(input)):
+        for y in range(height):
             vertical.append(input[y][x])
         lines.append(''.join(vertical))
     # Decreasing diagonals
@@ -43,5 +47,36 @@ def _extract_all_lines(input: Sequence[str]) -> Sequence[str]:
     return lines
 
 def count_xmas_words(input: Sequence[str]) -> int:
+    """Counts occurrences of the word XMAS in a grid.
+    The word XMAS may appear horizontally, vertically and diagonally.
+    It may appear reversed SMAX too.
+
+    Time complexity: O(n*nm) + O(m+n) = O(n^2m)
+    Space complexity: O(n+m)
+    """
     lines = _extract_all_lines(input)
-    return sum(line.count('XMAS') for line in lines) + sum(line.count('SAMX') for line in lines)
+    return sum(line.count('XMAS') for line in lines) + sum(line.count('SAMX') for line in lines)
+
+def count_xmas_shapes(input: Sequence[str]) -> int:
+    """Counts occurrences of the X-MAS shape in a grid.
+    There are 4 possible X-MAS shapes:
+    M S  S S  M M  S M
+     A    A    A    A
+    M S  M M  S S  S M
+
+    Time complexity: O(nm)
+    Space complexity: O(1)
+    """
+    width = len(input[0])
+    height = len(input)
+    counter = 0
+    for y in range(height):
+        for x in range(width):
+            # X-MAS shape must fit within bounds.
+            if 0 < y < height-1 and 0 < x < width-1 and input[y][x] == 'A':
+                found_shape = (((input[y-1][x-1] == 'M' and input[y+1][x+1] == 'S')
+                               or (input[y-1][x-1] == 'S' and input[y+1][x+1] == 'M'))
+                               and ((input[y+1][x-1] == 'M' and input[y-1][x+1] == 'S')
+                                or (input[y+1][x-1] == 'S' and input[y-1][x+1] == 'M')))
+                counter += 1 if found_shape else 0
+    return counter

aoc2024/test/day04/python/example3.txt

@@ -0,0 +1,3 @@
+M.S
+.A.
+M.S

aoc2024/test/day04/python/example4.txt

@@ -0,0 +1,10 @@
+.M.S......
+..A..MSMS.
+.M.S.MAA..
+..A.ASMSM.
+.M.S.M....
+..........
+S.S.S.S.S.
+.A.A.A.A..
+M.M.M.M.M.
+..........

aoc2024/test/day04/python/test_solution.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import unittest
 
-from aoc2024.src.day04.python.solution import count_xmas_words
+from aoc2024.src.day04.python.solution import count_xmas_words, count_xmas_shapes
 from common.python3.AdventOfCodeTestCase import AdventOfCodeTestCase
 
 
@@ -30,6 +30,15 @@ def test_part1_withExample2_counts(self):
   def test_part1_withPuzzleInput_counts(self):
     self.assertEqual(2603, count_xmas_words(self.input))
 
+  def test_part2_withExample_counts(self):
+    self.assertEqual(1, count_xmas_shapes(self.examples[2]))
+
+  def test_part2_withExample2_counts(self):
+    self.assertEqual(9, count_xmas_shapes(self.examples[3]))
+
+  def test_part2_withPuzzleInput_counts(self):
+    self.assertEqual(1965, count_xmas_shapes(self.input))
+
 
 if __name__ == '__main__':
   unittest.main()