From 2778ab290d0404edf587d80e80c15a29cf562aff Mon Sep 17 00:00:00 2001
From: Luc Rubio <luc@loociano.com>
Date: Wed, 4 Dec 2024 22:44:38 +0100
Subject: [PATCH] AoC 2024 day 4 part 2

---
 aoc2024/src/day04/python/solution.py       | 39 ++++++++++++++++++++--
 aoc2024/test/day04/python/example3.txt     |  3 ++
 aoc2024/test/day04/python/example4.txt     | 10 ++++++
 aoc2024/test/day04/python/test_solution.py | 11 +++++-
 4 files changed, 60 insertions(+), 3 deletions(-)
 create mode 100644 aoc2024/test/day04/python/example3.txt
 create mode 100644 aoc2024/test/day04/python/example4.txt

diff --git a/aoc2024/src/day04/python/solution.py b/aoc2024/src/day04/python/solution.py
index 47e2256..90c9026 100644
--- a/aoc2024/src/day04/python/solution.py
+++ b/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)
\ No newline at end of file
+    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
diff --git a/aoc2024/test/day04/python/example3.txt b/aoc2024/test/day04/python/example3.txt
new file mode 100644
index 0000000..8e6d4d8
--- /dev/null
+++ b/aoc2024/test/day04/python/example3.txt
@@ -0,0 +1,3 @@
+M.S
+.A.
+M.S
\ No newline at end of file
diff --git a/aoc2024/test/day04/python/example4.txt b/aoc2024/test/day04/python/example4.txt
new file mode 100644
index 0000000..a14c76e
--- /dev/null
+++ b/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.
+..........
\ No newline at end of file
diff --git a/aoc2024/test/day04/python/test_solution.py b/aoc2024/test/day04/python/test_solution.py
index ab15c22..c48ebc5 100644
--- a/aoc2024/test/day04/python/test_solution.py
+++ b/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()