solve day 6 part 2

day_06/program.py

@@ -37,19 +37,48 @@ def pos_invalid(row, col, max_row, max_col) -> bool:
     return True
 
 
-def get_visits(lab: list, g_row: int, g_col: int) -> int:
-    visits = {}
+def find_loop(lab: list, g_row: int, g_col: int, g_dir: int,
+              block_row: int, block_col: int) -> bool:
+    max_row = len(lab) - 1
+    max_col = len(lab[0]) - 1
+    if pos_invalid(block_row, block_col, max_row, max_col):
+        return False
+    test_visits = set()
+    while True:
+        if (g_row, g_col, g_dir) in test_visits:
+            return True
+        test_visits.add((g_row, g_col,g_dir))
+        new_row, new_col = next_position(g_row, g_col, g_dir)
+        if pos_invalid(new_row, new_col, max_row, max_col):
+            return False
+        elif lab[new_row][new_col] == "#" or (new_row == block_row and new_col == block_col):
+            g_dir = (g_dir + 1) % 4
+        else:
+            g_row = new_row
+            g_col = new_col
+
+
+def get_visits(lab: list, g_row: int, g_col: int) -> tuple:
+    visits = set()
+    obstructions = set()
+    obs_tested = set()
     g_dir = 0
+    start_row, start_col, start_dir = g_row, g_col, g_dir
     max_row = len(lab) - 1
     max_col = len(lab[0]) - 1
     while True:
+        visits.add((g_row, g_col))
         new_row, new_col = next_position(g_row, g_col, g_dir)
-        visits[g_row, g_col] = True
         if pos_invalid(new_row, new_col, max_row, max_col):
-            return len(visits)
+            return len(visits), len(obstructions)
         elif lab[new_row][new_col] == "#":
             g_dir = (g_dir + 1) % 4
         else:
+            block_row, block_col = next_position(g_row, g_col, g_dir)
+            if not (block_row, block_col) in obs_tested:
+                if find_loop(lab, start_row, start_col, start_dir, block_row, block_col):
+                    obstructions.add((block_row,block_col))
+                obs_tested.add((block_row, block_col))
             g_row = new_row
             g_col = new_col
 
@@ -58,8 +87,9 @@ def main():
     # lines = get_lines("sample-input.txt")
     lines = get_lines("input.txt")
     lab, g_row, g_col = get_lab(lines)
-    visits = get_visits(lab, g_row, g_col)
+    visits, obstructions = get_visits(lab, g_row, g_col)
     print("Part 1: Visited positions:", visits)
+    print("Part 2: Possible positions for obstructions:", obstructions)
 
 
 if __name__ == '__main__':