Solve 2024 day 6 part 2

2024/day-06/main.cpp

@@ -1,4 +1,5 @@
 #include <iostream>
+#include <set>
 #include <vector>
 
 using namespace std;
@@ -16,67 +17,104 @@ void readInput(vector<string> &mappedArea, pair<int, int> &startingPosition)
     }
 }
 
+bool isOutOfBounds(const vector<string> &mappedArea, const pair<int, int> &position)
+{
+    return (position.first < 0 || position.first >= mappedArea.size())
+           || (position.second < 0 || position.second >= mappedArea[0].size());
+}
+
+bool isObstacle(const vector<string> &mappedArea, const pair<int, int> &position)
+{
+    char cell = mappedArea[position.first][position.second];
+    return cell == '#' || cell == 'O';
+}
+
+enum Direction { UP = 0, RIGHT, DOWN, LEFT };
+
 bool takeStep(vector<string> &mappedArea,
               pair<int, int> &currentPosition,
-              int &currentOrientation,
-              int &distinctPositionCount)
+              Direction &currentDirection,
+              set<pair<int, int>> *distinctPositions = nullptr)
 {
-    pair<int, int> nextPosition = currentPosition;
+    const vector<pair<int, int>> directions = {{-1, 0}, {0, 1}, {1, 0}, {0, -1}};
+    pair<int, int> nextPosition = {currentPosition.first + directions[currentDirection].first,
+                                   currentPosition.second + directions[currentDirection].second};
 
-    switch (currentOrientation) {
-    case 0:
-        nextPosition.first--;
-        break;
-    case 1:
-        nextPosition.second++;
-        break;
-    case 2:
-        nextPosition.first++;
-        break;
-    case 3:
-        nextPosition.second--;
-        break;
-    default:
-        break;
-    }
-
-    if ((nextPosition.first >= mappedArea.size() || nextPosition.first < 0)
-        || (nextPosition.second >= mappedArea.front().size() || nextPosition.second < 0)) {
-        mappedArea[currentPosition.first][currentPosition.second] = '#';
+    if (isOutOfBounds(mappedArea, nextPosition)) {
         return false;
     }
 
-    if (mappedArea[nextPosition.first][nextPosition.second] == '#') {
-        currentOrientation = ++currentOrientation % 4;
-        return takeStep(mappedArea, currentPosition, currentOrientation, distinctPositionCount);
+    if (isObstacle(mappedArea, nextPosition)) {
+        currentDirection = static_cast<Direction>((currentDirection + 1) % 4);
+        return true;
     }
 
-    distinctPositionCount += !(mappedArea[nextPosition.first][nextPosition.second] == 'X');
+    if (distinctPositions && mappedArea[nextPosition.first][nextPosition.second] != 'X') {
+        distinctPositions->insert({nextPosition.first, nextPosition.second});
+    }
 
     mappedArea[currentPosition.first][currentPosition.second] = 'X';
     mappedArea[nextPosition.first][nextPosition.second] = '^';
-
     currentPosition = nextPosition;
     return true;
 }
 
+bool isLoopCausingObstruction(const pair<int, int> &obstructionPosition,
+                              vector<string> mappedArea,
+                              pair<int, int> currentPosition,
+                              Direction currentDirection)
+{
+    if (obstructionPosition == currentPosition) {
+        return false;
+    }
+
+    mappedArea[obstructionPosition.first][obstructionPosition.second] = 'O';
+
+    set<tuple<int, int, Direction>> states;
+    states.insert({currentPosition.first, currentPosition.second, currentDirection});
+
+    while (takeStep(mappedArea, currentPosition, currentDirection)) {
+        tuple<int, int, Direction> currentState = {currentPosition.first,
+                                                   currentPosition.second,
+                                                   currentDirection};
+
+        if (states.count(currentState)) {
+            return true;
+        }
+
+        states.insert(currentState);
+    }
+
+    return false;
+}
+
 int main()
 {
     vector<string> mappedArea;
-    pair<int, int> currentPosition;
-    readInput(mappedArea, currentPosition);
+    pair<int, int> startingPosition;
+    Direction startingDirection = UP;
+    readInput(mappedArea, startingPosition);
 
-    int currentOrientation = 0;
-    int distinctPositionCount = 1;
-    while (takeStep(mappedArea, currentPosition, currentOrientation, distinctPositionCount))
+    pair<int, int> currentPosition = startingPosition;
+    Direction currentDirection = startingDirection;
+    set<pair<int, int>> distinctPositions = {currentPosition};
+    while (takeStep(mappedArea, currentPosition, currentDirection, &distinctPositions))
         ;
+    int numberOfDistinctPositions = distinctPositions.size();
 
-    for (auto row : mappedArea) {
-        cout << row << endl;
+    int numberOfLoopCausingObstructions = 0;
+    for (const pair<int, int> &obstructionPosition : distinctPositions) {
+        numberOfLoopCausingObstructions += isLoopCausingObstruction(obstructionPosition,
+                                                                    mappedArea,
+                                                                    startingPosition,
+                                                                    startingDirection);
     }
-    cout << endl;
 
     // Part one
-    cout << "How many distinct positions will the guard visit before leaving the mapped area? "
-         << distinctPositionCount << endl;
+    cout << "Number of distinct positions visited before leaving: " << numberOfDistinctPositions
+         << endl;
+
+    // Part two
+    cout << "Number of potential obstructions to create a loop: " << numberOfLoopCausingObstructions
+         << endl;
 }