Python
Took me quite a while to figure out. Did part 3 using DFS initially, then optimized to use memoization.
# queue for BFS
from collections import deque
# possible moves for the dragon (knight-like moves)
DRAGON_MOVES = [(2, 1), (2, -1), (-2, 1), (-2, -1), (1, 2), (1, -2), (-1, 2), (-1, -2)]
# yields all possible moves for dragon at (i, j)
def get_next_moves(i, j, m, n):
for dx, dy in DRAGON_MOVES:
n_i, n_j = i + dx, j + dy
if 0 <= n_i < m and 0 <= n_j < n:
yield n_i, n_j
# BFS for given number of moves
def part1(data: str, moves: int = 4):
board = [list(row) for row in data.splitlines()]
m, n = len(board), len(board[0])
# initialize queue to dragon position
queue = deque()
for i in range(m):
for j in range(n):
if board[i][j] == 'D':
queue.append((i, j))
eaten = 0
visited = set()
# perform BFS for the given number of moves
# we do moves+1 because we only process a move after popping from queue
for _ in range(moves+1):
nq = len(queue)
for _ in range(nq):
i, j = queue.popleft()
if (i, j) in visited:
continue
visited.add((i, j))
# eat sheep if present
if board[i][j] == 'S':
eaten += 1
# add unvisited next moves to queue
for n_i, n_j in get_next_moves(i, j, m, n):
if (n_i, n_j) in visited:
continue
queue.append((n_i, n_j))
return eaten
# <assert snipped>
def part2(data: str, rounds: int = 20):
board = [list(row) for row in data.splitlines()]
m, n = len(board), len(board[0])
# initialize list to initial dragon location
dragon_locs = []
for i in range(m):
for j in range(n):
if board[i][j] == 'D':
dragon_locs.append((i, j))
eaten = 0
# simulate for given rounds
for _ in range(rounds):
# move dragon
# calculate all possible new positions for the dragon
dragon_moved_to = set()
for i, j in dragon_locs:
for n_i, n_j in get_next_moves(i, j, m, n):
dragon_moved_to.add((n_i, n_j))
# update dragon locations and eat any sheep present
dragon_locs = list(dragon_moved_to)
for i, j in dragon_locs:
if board[i][j] == 'S':
eaten += 1
board[i][j] = '.' # sheep is eaten
# move sheep
# process from bottom to top to avoid overlaps
for i in range(m-1, -1, -1):
for j in range(n):
if 'S' not in board[i][j]:
# no sheep here
continue
# move sheep out of current position
board[i][j] = board[i][j].replace('S', '')
if i == m-1:
# sheep escapes
continue
if board[i+1][j] == '#':
# sheep moves into hiding
board[i+1][j] = 'S#'
continue
if (i+1, j) in dragon_moved_to:
# sheep runs into dragon
eaten += 1
continue
# sheep moves down
board[i+1][j] = 'S'
return eaten
# <assert snipped>
# DP with memoization
def part3(data: str):
board_lines = data.splitlines()
m, n = len(board_lines), len(board_lines[0])
hideouts = set() # positions of hideouts
initial_sheep = [None] * n # initial positions of sheep in each column
dragon_loc = None # initial position of dragon
# iterate through board to find hideouts, sheep, and dragon
for i in range(m):
for j in range(n):
char = board_lines[i][j]
if char == '#':
hideouts.add((i, j))
elif char == 'D':
dragon_loc = (i, j)
elif char == 'S':
initial_sheep[j] = i
# convert to tuple for immutability and hashability
initial_sheep = tuple(initial_sheep)
# optional: calculate lost positions for each column
# lost_positions[j] is the row index where if the sheep reaches, can escape or stay hidden until escape (game over)
# used to quickly prune unwinnable states
lost_positions = [m] * n
for j in range(n):
for i in range(m-1, -1, -1):
if (i, j) not in hideouts:
break
lost_positions[j] = i
# memoization of state to number of paths
# state is a tuple of (player, dragon_loc, sheep_positions)
memo = {}
# do sheeps' turn
def solve_sheep(dragon_loc: tuple[int, int], sheep_positions: tuple[int|None, ...]) -> int:
# check for memoized result
state = ('S', dragon_loc, sheep_positions)
if state in memo:
return memo[state]
total_paths = 0
# used to check if any sheep can move
# this helps determine if the dragon gets another turn
can_move = False
for j, i in enumerate(sheep_positions):
if i is None:
# no sheep in this column
continue
next_i = i + 1
# Check collision with dragon
if (next_i, j) == dragon_loc and (next_i, j) not in hideouts:
# smart sheep avoids dragon
continue
# Check escape (entering safe zone)
if next_i == lost_positions[j]:
# this is a valid move, but its game over so we skip exploring further
can_move = True
continue
# this sheep will move down
can_move = True
# create new state with sheep moved down
new_sheep = list(sheep_positions)
new_sheep[j] = next_i
# continue solving for dragon's turn
total_paths += solve_dragon(dragon_loc, tuple(new_sheep))
# if no sheep could move, dragon gets another turn
if not can_move:
total_paths += solve_dragon(dragon_loc, sheep_positions)
# memoize result before returning
memo[state] = total_paths
return total_paths
# do dragon's turn
def solve_dragon(dragon_loc, sheep_positions):
# check for memoized result
state = ('D', dragon_loc, sheep_positions)
if state in memo:
return memo[state]
total_paths = 0
i, j = dragon_loc
# try all possible dragon moves
for n_i, n_j in get_next_moves(i, j, m, n):
# get where the sheep is in the column the dragon is moving to
sheep_loc_in_col = sheep_positions[n_j]
# if sheep is in the same row as the dragon and not in hideout, eat it
if sheep_loc_in_col == n_i and (n_i, n_j) not in hideouts:
sheep_count = sum(1 for x in sheep_positions if x is not None)
if sheep_count == 1:
# all sheep eaten, end of path
total_paths += 1
else:
# create new state with sheep eaten
new_sheep = list(sheep_positions)
new_sheep[n_j] = None
# continue solving for sheep's turn
total_paths += solve_sheep((n_i, n_j), tuple(new_sheep))
else:
# Empty, hideout, or sheep hidden in hideout
# continue solving for sheep's turn with dragon moved
total_paths += solve_sheep((n_i, n_j), sheep_positions)
# memoize result before returning
memo[state] = total_paths
return total_paths
# start solving with sheep's turn
return solve_sheep(dragon_loc, initial_sheep)
# <asserts snipped>