adventofcode/2024/06/part1.py

import numpy as np
from enum import Enum
import copy


verbose = False

class Map:
	"""
	Handles a map with obstacles and how a guard walks through it.
	"""

	directions = ["^", ">", "v", "<"]
	step_per_direction = [(-1, 0), (0, 1), (1, 0), (0, -1)]

	class Tiles(Enum):
		BORDER = "+"
		FREE = "."
		OBSTACLE = "#"

	def __init__(self):
		self.trace = []
		self.pos = None
		self.direction = None

	def next_pos(self):
		"""
		Returns next position given the state of this Map.
		"""

		next_pos = tuple([self.pos[i] + self.step_per_direction[self.directions.index(self.direction)][i] for i in range(len(self.pos))])
		if verbose:
			print("next_pos", next_pos)
		return next_pos

	def pos_in_map(self, pos):
		"""
		Returns whether pos is in the map.
		"""

		if verbose:
			print("pos_in_map", pos) 
		return pos[0] >= 0 and pos[0] < self.map.shape[0] \
		   and pos[1] >= 0 and pos[1] < self.map.shape[1]

	def look_ahead(self):
		next_step = self.next_pos()
		if verbose:
			print("look_ahead.next_step", next_step)

		if not self.pos_in_map(next_step):
			return self.Tiles.BORDER
		if self.map[next_step] == "#":
			return self.Tiles.OBSTACLE
		if self.map[next_step] == ".":
			return self.Tiles.FREE

	def next_direction(self, direction: str):
		"""
		Return the direction turning right given direction.
		"""

		assert direction in self.directions, "Invalid direction given"

		if self.directions[-1] == direction:
			return self.directions[0]
		else:
			return self.directions[self.directions.index(direction) + 1]

	def seek_guard_in_map(self):
		"""
		Looks for the guard in the map and returns its position.

		Will raise an exception if multiple or none are found.
		"""
		
		guard_locations = []
		for direction in self.directions:
			guard_locations.extend(list(zip(*np.where(self.map == direction))))
		return guard_locations
	
	def load_map(self, filename="input"):
		"""
		Reads a map from a file without verifying its contents.
		"""

		with open(filename, "r") as fp:
			data = fp.read().splitlines()
			chararray = np.array(data, dtype=str)\
				.view("U1")\
				.reshape((len(data), -1))
		self.map = chararray

		# Find the location of the guard in the map.
		guards = self.seek_guard_in_map()

		if verbose:
			print(guards)
		assert len(guards) == 1, "There should only be one guard in the map."

		self.pos = guards[0]
		self.trace.append(self.pos)
		self.direction = self.map[self.pos]
		self.map[self.pos] = self.Tiles.FREE.value

	class LoopException(Exception):
		"""
		Gets thrown when the map results in the guard walking in an
		infinite loop.
		"""

		pass


	def step(self):
		"""
		Take a step, turn or stop.
		"""

		next_tile = self.look_ahead()
		if verbose:
			print("step.next_tile", next_tile)

		if self.is_stuck_in_loop():
			raise self.LoopException()

		if next_tile == self.Tiles.FREE:
			# TODO: Really take a step! Do not only turn.
			self.pos = self.next_pos()
			self.trace.append(self.pos)
			return True
		elif next_tile == self.Tiles.OBSTACLE:
			self.direction = self.next_direction(self.direction)
			return True
		elif next_tile == self.Tiles.BORDER:
			return False

		assert False, "Should not come here :')"	

	def show(self):
		"""Prints the current map to stdout."""

		for i in range(len(self.map)):
			for j in range(len(self.map[0])):
				if self.pos == (i, j):
					print(self.direction, end="")
				else:
					print(self.map[i][j], end="")
			print()

	def is_stuck_in_loop(self):
		"""
		Returns whether the guard is stuck in loop.

		This only happens when enough steps are taken. Currently, it
		checks whether the current position is already part of the trace
		at least four times, so that it has travelled the position at
		least twice from at least one direction.
		"""

		return self.trace.count(self.pos) > 4

	def copy(self):
		if verbose:
			print(f"Deep copying")
		return copy.deepcopy(self)


if __name__ == "__main__":
	m = Map()
	#m.load_map("testinput")
	m.load_map("input")

	loop_obstacles = []

	while m.step():
		if verbose:
			print(m.pos, m.direction)

		# Create a copy
		forked_map = m.copy()
		forked_map.map[forked_map.next_pos()] = Map.Tiles.OBSTACLE.value
		print(f"Forking at {forked_map.next_pos()}... ", end="")
		try:
			while forked_map.step():
				continue
		except Map.LoopException:
			print("loop.")
			loop_obstacles.append(m.next_pos())
		else:
			print("no loop")
		#m.show()
	#print(len(m.trace), m.trace)
	#print(len(set(m.trace)))
	print(len(loop_obstacles))