Dazbo's Advent of Code solutions, written in Python
dataclasslist comprehensionsetsassertionUnit TestingWorking with ImagesVisualisations with MatplotlibTiming with tqdm
This one was fun. Solving the problem took me an hour or so. But then I spent far too long creating a visualisation!
The goal is model the behaviour of a strange rope. The rope has a knot at each end: the head and the tail. At the beginning, both the head and tail occupy the same space.
The head of the rope is moved according to a set of instructions that look like this:
R 5
U 8
L 8
D 3
R 17
D 10
L 25
U 20
We’re told that whenever we move the head, the tail must move in order to remain adjacent to the head.
How many positions does the tail of the rope visit at least once?
My strategy:
Point
dataclass to store locations.
Point (0,0)
.set
. This is how we record locations that have been visited at least once.- Then process each movement instruction:
First, the Point
dataclass:
@dataclass(frozen=True)
class Point:
""" Class for storing a point x,y coordinate """
x: int
y: int
# create a list of (x,y) vectors that surround and include this point
WITHIN_ONE = [(dx,dy) for dx in range(-1, 2) for dy in range(-1, 2)]
def __add__(self, other):
return Point(self.x + other.x, self.y + other.y)
def __sub__(self, other):
return Point(self.x - other.x, self.y - other.y)
Notes on this class:
__add__()
and __sub__()
methods, thus allowing points to be added or subtracted using standard operators, i.e. using +
and -
, respectively.[(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0), (0, 1), (1, -1), (1, 0), (1, 1)]
Then we’ve created VECTOR dict
, to convert our input instructions to points/vectors:
VECTORS = {
'U': Point(0, 1),
'R': Point(1, 0),
'D': Point(0, -1),
'L': Point(-1, 0)
}
Finally, I create a RopeSim class
:
class RopeSim():
""" Simulates a rope with a number of knots. We move the head according to a set of instructions.
Here we model the movement of the knots behind the head, according to the rules specified. """
def __init__(self, motions: list[tuple[str, int]], num_knots: int) -> None:
""" Expects a list of instructions in the format:
[['R', 5], ['U', 8], ...]
Models rope with num_knots. The first is the head, and the last is the tail. """
self._instructions = motions
self._num_knots = num_knots
self._knots = [Point(0,0) for _ in range(self._num_knots)]
@staticmethod
def _get_next_move(vector: Point) -> Point:
x_move = y_move = 0
move_x = move_y = False
if vector.y == 0: # we only need to move left or right
move_x = True
elif vector.x == 0: # we only need to move up or down
move_y = True
else: # we need to move diagonally
assert vector.x != 0 and vector.y != 0, "We must move diagonally"
move_x = move_y = True
if move_x:
x_move = 1 if vector.x > 0 else -1
if move_y:
y_move = 1 if vector.y > 0 else -1
return Point(x_move, y_move)
def pull_rope(self) -> set[Point]:
""" Simulate the rope knot movemens, according to the rules given. """
visited_locations: set[Point] = set()
visited_locations.add(self._knots[-1]) # track the tail
for direction, mag in self._instructions: # read char by char
for _ in range(mag): # move one step at a time
# print(f"Tail: {knots[-1]}; unique positions: {len(visited_locations)}")
self._knots[0] += VECTORS[direction] # move the head
for i in range(1, len(self._knots)): # move the tail
vector = self._knots[i-1] - self._knots[i]
if vector in [Point(x,y) for (x,y) in Point.WITHIN_ONE]:
continue # don't need to move
else:
self._knots[i] = self._knots[i] + RopeSim._get_next_move(vector)
visited_locations.add(self._knots[-1])
return visited_locations
I didn’t really need to use a class for this. But later, I wanted to add a visualisation, and doing this with a class just made life easier.
__init__()
method:
2
knots: the head and tail.Point
of 0,0
._get_next_move(vector: Point)
method:
x
and y
components of the vector will never exceed 1.pull_rope()
method:
visited set
to store every location that tail has been.WITHIN_ONE
list. If it is, then the tail is already adjacent to (or in the same spot as) the head._get_next_move()
method to determine the movement required for the tail.Now all we need to do is read our data, pass it to our RopeSim
, and execute the pull_rope()
method:
def main():
with open(INPUT_FILE, mode="rt") as f:
# convert to list of (direction, magnitude)
data = [(d, int(v)) for d, v in [instruction.split() for instruction in f.read().splitlines()]]
rope_sim = RopeSim(data, 2)
visited_locations = rope_sim.pull_rope()
print(len(visited_locations))
Now we’re told the rope has mulitple knots. Each knot behaves like the tail to the knot in front of it.
Simulate your complete series of motions on a larger rope with ten knots. How many positions does the tail of the rope visit at least once?
Marvellous!! I don’t need to change anything! I just need to pass in 10
knots, instead of 2
.
rope_sim = RopeSim(data, 10)
visited_locations = rope_sim.pull_rope()
print(len(visited_locations))
The final code looks like this:
from dataclasses import dataclass
from pathlib import Path
import time
SCRIPT_DIR = Path(__file__).parent
# INPUT_FILE = Path(SCRIPT_DIR, "input/sample_input.txt")
INPUT_FILE = Path(SCRIPT_DIR, "input/input.txt")
@dataclass(frozen=True)
class Point:
""" Class for storing a point x,y coordinate """
x: int
y: int
# create a list of (x,y) vectors that sorround and include this point
WITHIN_ONE = [(dx,dy) for dx in range(-1, 2) for dy in range(-1, 2)]
def __add__(self, other):
return Point(self.x + other.x, self.y + other.y)
def __sub__(self, other):
return Point(self.x - other.x, self.y - other.y)
VECTORS = {
'U': Point(0, 1),
'R': Point(1, 0),
'D': Point(0, -1),
'L': Point(-1, 0)
}
def main():
with open(INPUT_FILE, mode="rt") as f:
# convert to list of (direction, magnitude)
data = [(d, int(v)) for d, v in [instruction.split() for instruction in f.read().splitlines()]]
answers = []
for num_knots in (2, 10)
rope_sim = RopeSim(data, num_knots)
visited_locations = rope_sim.pull_rope()
answers.append(len(visited_locations))
print(answers)
class RopeSim():
""" Simulates a rope with a number of knots. We move the head according to a set of instructions.
Here we model the movement of the knots behind the head, according to the rules specified. """
def __init__(self, motions: list[tuple[str, int]], num_knots: int) -> None:
""" Expects a list of instructions in the format:
[['R', 5], ['U', 8], ...]
Models rope with num_knots. The first is the head, and the last is the tail. """
self._instructions = motions
self._num_knots = num_knots
self._knots = [Point(0,0) for _ in range(self._num_knots)]
@staticmethod
def _get_next_move(vector: Point) -> Point:
x_move = y_move = 0
move_x = move_y = False
if vector.y == 0: # we only need to move left or right
move_x = True
elif vector.x == 0: # we only need to move up or down
move_y = True
else: # we need to move diagonally
assert vector.x != 0 and vector.y != 0, "We must move diagonally"
move_x = move_y = True
if move_x:
x_move = 1 if vector.x > 0 else -1
if move_y:
y_move = 1 if vector.y > 0 else -1
return Point(x_move, y_move)
def pull_rope(self) -> set[Point]:
""" Simulate the rope knot movemens, according to the rules given. """
visited_locations: set[Point] = set()
visited_locations.add(self._knots[-1]) # track the tail
for direction, mag in self._instructions: # read char by char
for _ in range(mag): # move one step at a time
# print(f"Tail: {knots[-1]}; unique positions: {len(visited_locations)}")
self._knots[0] += VECTORS[direction] # move the head
for i in range(1, len(self._knots)): # move the tail
vector = self._knots[i-1] - self._knots[i]
if vector in [Point(x,y) for (x,y) in Point.WITHIN_ONE]:
continue # don't need to move
else:
self._knots[i] = self._knots[i] + RopeSim._get_next_move(vector)
visited_locations.add(self._knots[-1])
return visited_locations
if __name__ == "__main__":
t1 = time.perf_counter()
main()
t2 = time.perf_counter()
print(f"Execution time: {t2 - t1:0.4f} seconds")
The output looks like this:
6376
2607
Execution time: 0.4464 seconds
This problem seems like a prime candidate for a cool visualisation. So I’ve had a go at creating an animation.
First, we need some extra imports:
from io import BytesIO
from pathlib import Path
import time
import imageio as iio
from tqdm import tqdm
from matplotlib import pyplot as plt
from matplotlib.markers import MarkerStyle
from matplotlib.ticker import MaxNLocator
Then I create an Animator
class, which just happens to be Context Manager. I’ll talk about that another time!
class Animator():
""" Creates an animation file of specified target size.
Designed to be used as Context Manager. E.g.
with Animator(file=Path("path/to/file""), fps=num) as animator:
# code
"""
def __enter__(self):
""" Required for ContextManager implementation. """
if self._enabled:
self._create_path()
return self # so the as <name> returns an object
def __exit__(self, exc_type, exc_val, exc_tb):
""" Required for ContextManager implementation. """
if self._enabled:
self._save_anim()
def __init__(self, file: Path, fps: int, loop=1, enabled=True) -> None:
""" Create an Animator. Suggest the file should be a .gif.
Set frames per second (fps).
Set loop to 0, to loop indefinitely. Default is 1. """
self._enabled = enabled
self._outputfile = file
self._frames = [] # can be in-memory BytesIO objects, or files
self._fps = fps
self._loop = loop
@property
def enabled(self):
return self._enabled
@enabled.setter
def enabled(self, value: bool):
self._enabled = value
def _create_path(self):
if self._enabled:
dir_path = Path(self._outputfile).parent
if not Path.exists(dir_path):
Path.mkdir(dir_path)
def _save_anim(self):
""" Takes animation frames, and converts to a single animated file. """
with iio.get_writer(self._outputfile, mode='I', fps=self._fps, loop=self._loop) as writer:
for frame in tqdm(self._frames):
image = iio.imread(frame)
writer.append_data(image)
print(f"Animation saved to {self._outputfile}")
def add_frame(self, frame):
""" Add a frame to the animation.
The frame can be in the form of a BytesIO object, or a file Path
"""
self._frames.append(frame)
I won’t go into this class in detail. But a few notes…
__enter__()
and __exit__()
methods allow me to use this class as a Context Manager, i.e. so that I can say with Animator...
, much like we do when we say with open(...)
.__init__()
method accepts a parameter that can be used to DISABLE this animator. This is handy, since rendering animations is time consuming.add_frame()
method is intended to be called by our application code, whenever we have a frame - such as Maplotlib plot - that we’re ready to add to our animation._save_anim()
writes our frames to a file on the disk. This method gets called automatically at the end of our with
context block, assuming the Animator
has not been disabled. We also use tqdm
to show a progress bar on the screen as we build the animation.Now we need to actually generate our visual frames. To do that, we add the following two methods to our RopeSim
class:
def _init_plt(self):
""" Generate a Figure and Axes objects which are reused. """
my_dpi = 120
figure, axes = plt.subplots(figsize=(1024/my_dpi, 768/my_dpi), dpi=my_dpi, facecolor="white") # set size in pixels
axes.set_aspect('equal') # set x and y to equal aspect
axes.set_facecolor('xkcd:black')
return figure, axes
def _render_frame(self, visited: set[Point], iteration: int=0):
""" Only renders an animation frame if we've attached an enabled Animator """
fig, axes = self._plt_info
axes.clear()
# The grid will grow as the rope heads moves around
max_x = max(self._all_head_locations, key=lambda point: point.x).x
min_x = min(self._all_head_locations, key=lambda point: point.x).x
max_y = max(self._all_head_locations, key=lambda point: point.y).y
min_y = min(self._all_head_locations, key=lambda point: point.y).y
axes.set_xlim(min_x - 2, max_x + 2)
axes.set_ylim(min_y - 2, max_y + 2)
# dynamically compute the marker size
fig.canvas.draw()
factor = 40 # Smaller factor means smaller markers
mkr_size = int((axes.get_window_extent().width / (max_x-min_x+1) * (factor/fig.dpi)) ** 2)
# make sure the ticks have integer values
axes.xaxis.set_major_locator(MaxNLocator(integer=True))
head = self._knots[0]
tail = self._knots[-1]
others_knots = self._knots[1:-1]
visited_x = [point.x for point in visited if point != tail]
visited_y = [point.y for point in visited if point != tail]
for knot in others_knots:
axes.scatter(knot.x, knot.y, marker=MarkerStyle("."), s=mkr_size/2, color="white")
axes.scatter(head.x, head.y, marker=MarkerStyle("."), s=mkr_size, color="red")
axes.scatter(visited_x, visited_y, marker=MarkerStyle("x"), s=mkr_size/3, color="grey")
axes.scatter(tail.x, tail.y, marker=MarkerStyle("*"), s=mkr_size/2, color="yellow")
axes.set_title(f"Iteration: {iteration}; tail has visited {len(visited)} locations")
# save the plot as a frame; store the frame in-memory, using a BytesIO buffer
frame = BytesIO()
plt.savefig(frame, format='png') # save to memory, rather than file
self._animator.add_frame(frame)
The _render_frame()
method uses Matplotlib to render a scatter plot, using different colours for the head, the tail, the knots in between, and the visited set. At the end of the method, we save the current plot to a BytesIO
object in memory, and then call our Animator add_frame()
method to add it to the list of frames that will ultimately be converted to a file-based animation.
We also need to set the animator and initialise our Matplotlib plot, in our RopeSim __init__()
method:
self._animator: Animator = animator
self._plt_info = self._init_plt() # contains (figure, axes)
Finally, we need to call _render_frame()
whenever we move a knot. So our pull_rope()
method now looks like this:
def pull_rope(self) -> set[Point]:
""" Simulate the rope knot movemens, according to the rules given. """
visited_locations: set[Point] = set()
visited_locations.add(self._knots[-1]) # track the tail
step = 0
for direction, mag in self._instructions: # read char by char
for _ in range(mag): # move one step at a time
# print(f"Tail: {knots[-1]}; unique positions: {len(visited_locations)}")
self._knots[0] += VECTORS[direction] # move the head
self._all_head_locations.add(self._knots[0])
for i in range(1, len(self._knots)): # move the tail
step += 1
vector = self._knots[i-1] - self._knots[i]
if vector in [Point(x,y) for (x,y) in Point.WITHIN_ONE]:
continue # don't need to move
else:
self._knots[i] = self._knots[i] + RopeSim._get_next_move(vector)
visited_locations.add(self._knots[-1])
if self._animator and self._animator.enabled:
self._render_frame(visited_locations, step)
return visited_locations
Here’s what it looks like:
So, my final code, including visualisation, looks like this:
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
import time
import imageio as iio
from tqdm import tqdm
from matplotlib import pyplot as plt
from matplotlib.markers import MarkerStyle
from matplotlib.ticker import MaxNLocator
SCRIPT_DIR = Path(__file__).parent
INPUT_FILE = Path(SCRIPT_DIR, "input/sample_input.txt")
# INPUT_FILE = Path(SCRIPT_DIR, "input/input.txt")
# MAKE SURE YOU DISABLE IF RUNNING WITH THE REAL DATA. IT TAKES TOO LONG!!!
ENABLE_ANIMATIONS = True
OUTPUT_FOLDER = Path(SCRIPT_DIR, "output/")
@dataclass(frozen=True)
class Point:
""" Class for storing a point x,y coordinate """
x: int
y: int
# create a list of (x,y) vectors that sorround and include this point
WITHIN_ONE = [(dx,dy) for dx in range(-1, 2) for dy in range(-1, 2)]
def __add__(self, other):
return Point(self.x + other.x, self.y + other.y)
def __sub__(self, other):
return Point(self.x - other.x, self.y - other.y)
class Animator():
""" Creates an animation file of specified target size.
Designed to be used as Context Manager. E.g.
with Animator(file=Path("path/to/file""), fps=num) as animator:
# code
"""
def __enter__(self):
""" Required for ContextManager implementation. """
if self._enabled:
self._create_path()
return self # so the as <name> returns an object
def __exit__(self, exc_type, exc_val, exc_tb):
""" Required for ContextManager implementation. """
if self._enabled:
self._save_anim()
def __init__(self, file: Path, fps: int, loop=1, enabled=True) -> None:
""" Create an Animator. Suggest the file should be a .gif.
Set frames per second (fps).
Set loop to 0, to loop indefinitely. Default is 1. """
self._enabled = enabled
self._outputfile = file
self._frames = [] # can be in-memory BytesIO objects, or files
self._fps = fps
self._loop = loop
@property
def enabled(self):
return self._enabled
@enabled.setter
def enabled(self, value: bool):
self._enabled = value
def _create_path(self):
if self._enabled:
dir_path = Path(self._outputfile).parent
if not Path.exists(dir_path):
Path.mkdir(dir_path)
def _save_anim(self):
""" Takes animation frames, and converts to a single animated file. """
with iio.get_writer(self._outputfile, mode='I', fps=self._fps, loop=self._loop) as writer:
for frame in tqdm(self._frames):
image = iio.imread(frame)
writer.append_data(image)
print(f"Animation saved to {self._outputfile}")
def add_frame(self, frame):
""" Add a frame to the animation.
The frame can be in the form of a BytesIO object, or a file Path
"""
self._frames.append(frame)
VECTORS = {
'U': Point(0, 1),
'R': Point(1, 0),
'D': Point(0, -1),
'L': Point(-1, 0)
}
def main():
with open(INPUT_FILE, mode="rt") as f:
# convert to list of (direction, magnitude)
data = [(d, int(v)) for d, v in [instruction.split() for instruction in f.read().splitlines()]]
answers = []
with Animator(file=Path(OUTPUT_FOLDER, "rope_bridge_pt1.gif"), fps=10, enabled=ENABLE_ANIMATIONS) as animator:
rope_sim = RopeSim(data, 2, animator=animator)
visited_locations = rope_sim.pull_rope()
answers.append(len(visited_locations))
with Animator(file=Path(OUTPUT_FOLDER, "rope_bridge_pt2.gif"), fps=20, enabled=ENABLE_ANIMATIONS) as animator:
rope_sim = RopeSim(data, 10, animator=animator)
visited_locations = rope_sim.pull_rope()
answers.append(len(visited_locations))
print(answers)
class RopeSim():
""" Simulates a rope with a number of knots. We move the head according to a set of instructions.
Here we model the movement of the knots behind the head, according to the rules specified. """
def __init__(self, motions: list[tuple[str, int]], num_knots: int, animator=None) -> None:
""" Expects a list of instructions in the format:
[['R', 5], ['U', 8], ...]
Models rope with num_knots. The first is the head, and the last is the tail. """
self._instructions = motions
self._num_knots = num_knots
self._knots = [Point(0,0) for _ in range(self._num_knots)]
self._all_head_locations: set[Point] = set() # for rendering the vis
self._animator: Animator = animator
self._plt_info = self._init_plt() # contains (figure, axes)
@staticmethod
def _get_next_move(vector: Point) -> Point:
x_move = y_move = 0
move_x = move_y = False
if vector.y == 0: # we only need to move left or right
move_x = True
elif vector.x == 0: # we only need to move up or down
move_y = True
else: # we need to move diagonally
assert vector.x != 0 and vector.y != 0, "We must move diagonally"
move_x = move_y = True
if move_x:
x_move = 1 if vector.x > 0 else -1
if move_y:
y_move = 1 if vector.y > 0 else -1
return Point(x_move, y_move)
def pull_rope(self) -> set[Point]:
""" Simulate the rope knot movemens, according to the rules given. """
visited_locations: set[Point] = set()
visited_locations.add(self._knots[-1]) # track the tail
step = 0
for direction, mag in self._instructions: # read char by char
for _ in range(mag): # move one step at a time
# print(f"Tail: {knots[-1]}; unique positions: {len(visited_locations)}")
self._knots[0] += VECTORS[direction] # move the head
self._all_head_locations.add(self._knots[0])
for i in range(1, len(self._knots)): # move the tail
step += 1
vector = self._knots[i-1] - self._knots[i]
if vector in [Point(x,y) for (x,y) in Point.WITHIN_ONE]:
continue # don't need to move
else:
self._knots[i] = self._knots[i] + RopeSim._get_next_move(vector)
visited_locations.add(self._knots[-1])
if self._animator and self._animator.enabled:
self._render_frame(visited_locations, step)
return visited_locations
def _init_plt(self):
""" Generate a Figure and Axes objects which are reused. """
my_dpi = 120
figure, axes = plt.subplots(figsize=(1024/my_dpi, 768/my_dpi), dpi=my_dpi, facecolor="white") # set size in pixels
axes.set_aspect('equal') # set x and y to equal aspect
axes.set_facecolor('xkcd:black')
return figure, axes
def _render_frame(self, visited: set[Point], iteration: int=0):
""" Only renders an animation frame if we've attached an enabled Animator """
fig, axes = self._plt_info
axes.clear()
# The grid will grow as the rope heads moves around
max_x = max(self._all_head_locations, key=lambda point: point.x).x
min_x = min(self._all_head_locations, key=lambda point: point.x).x
max_y = max(self._all_head_locations, key=lambda point: point.y).y
min_y = min(self._all_head_locations, key=lambda point: point.y).y
axes.set_xlim(min_x - 2, max_x + 2)
axes.set_ylim(min_y - 2, max_y + 2)
# dynamically compute the marker size
fig.canvas.draw()
factor = 40 # Smaller factor means smaller markers
mkr_size = int((axes.get_window_extent().width / (max_x-min_x+1) * (factor/fig.dpi)) ** 2)
# make sure the ticks have integer values
axes.xaxis.set_major_locator(MaxNLocator(integer=True))
head = self._knots[0]
tail = self._knots[-1]
others_knots = self._knots[1:-1]
visited_x = [point.x for point in visited if point != tail]
visited_y = [point.y for point in visited if point != tail]
for knot in others_knots:
axes.scatter(knot.x, knot.y, marker=MarkerStyle("."), s=mkr_size/2, color="white")
axes.scatter(head.x, head.y, marker=MarkerStyle("."), s=mkr_size, color="red")
axes.scatter(visited_x, visited_y, marker=MarkerStyle("x"), s=mkr_size/3, color="grey")
axes.scatter(tail.x, tail.y, marker=MarkerStyle("*"), s=mkr_size/2, color="yellow")
axes.set_title(f"Iteration: {iteration}; tail has visited {len(visited)} locations")
# save the plot as a frame; store the frame in-memory, using a BytesIO buffer
frame = BytesIO()
plt.savefig(frame, format='png') # save to memory, rather than file
self._animator.add_frame(frame)
if __name__ == "__main__":
t1 = time.perf_counter()
main()
t2 = time.perf_counter()
print(f"Execution time: {t2 - t1:0.4f} seconds")
I’ve added a unit test to this application. I use this test to check that my application gives the right answers for the sample input.
Here, I’ve created a separate py script - called test_rope_bridge.py
- in the same package folder:
from pathlib import Path
import unittest
import rope_bridge
SAMPLE_INPUT_FILE = Path(Path(__file__).parent, "input/sample_input.txt")
class TestRopeBridge(unittest.TestCase):
""" Set up data using the sample input.
Then run two tests, asserting the correct length of the returned lists. """
def setUp(self):
# load the data
with open(SAMPLE_INPUT_FILE, mode="rt") as f:
self.data = [(d, int(v)) for d, v in [instruction.split() for instruction in f.read().splitlines()]]
def test_part_1(self):
expected = 88
rope_sim = rope_bridge.RopeSim(self.data, 2)
self.assertEqual(len(rope_sim.pull_rope()), expected)
def test_part_2(self):
expected = 36
rope_sim = rope_bridge.RopeSim(self.data, 10)
self.assertEqual(len(rope_sim.pull_rope()), expected)
if __name__ == "__main__":
unittest.main()
This works as follows:
TestRopeBridge
class, by extending unittest.TestCase
.setUp()
method, I read in the sample input, and store it as self.data
.RopeSim
class.