Python
# get all valid neighbors of a cell in a grid
DELTAS = [(-1, 0), (1, 0), (0, -1), (0, 1)]
def getNeighbors(i, j, m, n):
for di, dj in DELTAS:
ni, nj = i + di, j + dj
if 0 <= ni < m and 0 <= nj < n:
yield ni, nj
# constant for marking permanently burnt cells
PERMA_BURNT = -1
# DFS to burn the grid from a list of starting points
# If perma_burn is True, cells that are burnt will be modified to PERMA_BURNT
# returns the set of all burnt cells
def dfs_burn(grid: list[list[int]], start: list[tuple[int, int]], perma_burn=False) -> set[tuple[int, int]]:
m, n = len(grid), len(grid[0])
ignited = start
burnt = set()
while ignited:
i, j = ignited.pop()
if (i, j) in burnt:
continue
burnt.add((i, j))
for ni, nj in getNeighbors(i, j, m, n):
if (ni, nj) in burnt or grid[ni][nj] == PERMA_BURNT:
continue
if grid[ni][nj] <= grid[i][j]:
ignited.append((ni, nj))
# mark as permanently burnt (if requested)
if perma_burn: grid[i][j] = PERMA_BURNT
return burnt
# Part 1: DFS burn from single source (0, 0)
def part1(data: str) -> int:
grid = [[int(c) for c in row] for row in data.splitlines()]
return len(dfs_burn(grid, [(0, 0)]))
assert part1("""989601
857782
746543
766789""") == 16
# Part 2: DFS burn from two sources (0,0) and (m-1,n-1)
def part2(data: str) -> int:
grid = [[int(c) for c in row] for row in data.splitlines()]
m, n = len(grid), len(grid[0])
return len(dfs_burn(grid, [(0, 0), (m - 1, n - 1)]))
assert part2("""9589233445
9679121695
8469121876
8352919876
7342914327
7234193437
6789193538
6781219648
5691219769
5443329859""") == 58
from copy import deepcopy
# Part 3: Greedy selection of three best burn points
def part3(data: str) -> int:
grid = [[int(c) for c in row] for row in data.splitlines()]
m, n = len(grid), len(grid[0])
# keep original grid for final burn count
og_grid = deepcopy(grid)
# generate the list of best burn candidates
# note that the cell with the max value is not necessarily the single best burn point
# this will help us skip the cells that are clearly suboptimal
candidates = [(grid[i][j], i, j) for i in range(m) for j in range(n)]
candidates.sort(reverse=True)
# best three burn points (separately chosen)
best_three = []
for _ in range(3):
# set of all cells burnt in this iteration
burnt = set()
# track the best burn location in this iteration
max_burnt = 0
max_burn_loc = None
for _, i, j in candidates:
# skip candidates that are already burnt by a previous burn point
# this is because a previous burn point will contain all the burns of this candidate
# plus possibly more
if (i, j) in burnt:
continue
# burn (impermanently) from this candidate
burns = dfs_burn(grid, [(i, j)])
if len(burns) > max_burnt:
max_burnt = len(burns)
max_burn_loc = (i, j)
# record newly burnt cells to skip in future candidates
burnt |= burns
assert max_burn_loc is not None, "Should have found a max burn location"
# record and permanently burn from the best location found
dfs_burn(grid, [max_burn_loc], perma_burn=True)
best_three.append(max_burn_loc)
# return total burnt cells from all three best burn points simultaneously
return len(dfs_burn(og_grid, best_three))
assert (t := part3("""5411
3362
5235
3112""")) == 14, f"Expected: 14, Actual: {t}"
assert (t := part3("""41951111131882511179
32112222211518122215
31223333322115122219
31234444432147511128
91223333322176121892
61112222211166431583
14661111166111111746
11111119142122222177
41222118881233333219
71222127839122222196
56111126279711111517""")) == 136, f"Expected: 136, Actual: {t}"