Co-authored-by: Mikael CAPELLE <mikael.capelle@thalesaleniaspace.com> Co-authored-by: Mikaël Capelle <capelle.mikael@gmail.com> Reviewed-on: #3
60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
import itertools as it
|
|
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
|
|
from ..base import BaseSolver
|
|
|
|
|
|
def move(head: tuple[int, int], command: str) -> tuple[int, int]:
|
|
h_col, h_row = head
|
|
|
|
if command == "L":
|
|
head = (h_col - 1, h_row)
|
|
elif command == "R":
|
|
head = (h_col + 1, h_row)
|
|
elif command == "U":
|
|
head = (h_col, h_row + 1)
|
|
elif command == "D":
|
|
head = (h_col, h_row - 1)
|
|
|
|
return head
|
|
|
|
|
|
def follow(head: tuple[int, int], tail: tuple[int, int]) -> tuple[int, int]:
|
|
h_col, h_row = head
|
|
t_col, t_row = tail
|
|
|
|
if abs(t_col - h_col) <= 1 and abs(t_row - h_row) <= 1:
|
|
return tail
|
|
|
|
return t_col + np.sign(h_col - t_col), t_row + np.sign(h_row - t_row)
|
|
|
|
|
|
def run(commands: list[str], n_blocks: int) -> list[tuple[int, int]]:
|
|
blocks: list[tuple[int, int]] = [(0, 0) for _ in range(n_blocks)]
|
|
visited = [blocks[-1]]
|
|
|
|
for command in commands:
|
|
blocks[0] = move(blocks[0], command)
|
|
|
|
for i in range(0, n_blocks - 1):
|
|
blocks[i + 1] = follow(blocks[i], blocks[i + 1])
|
|
|
|
visited.append(blocks[-1])
|
|
|
|
return visited
|
|
|
|
|
|
class Solver(BaseSolver):
|
|
def solve(self, input: str) -> Iterator[Any]:
|
|
lines = [line.strip() for line in input.splitlines()]
|
|
|
|
# flatten the commands
|
|
commands = list(
|
|
it.chain(*(p[0] * int(p[1]) for line in lines if (p := line.split())))
|
|
)
|
|
|
|
yield len(set(run(commands, n_blocks=2)))
|
|
yield len(set(run(commands, n_blocks=10)))
|