import os, argparse import numpy as np debug_print = print debug_print = lambda x: x def solve_task(lines): s = Snake(9) s.print() for line in lines: s.do(line) #s.print() print(f"Tail visited {s.tail_visited()} spaces") class Snake: def __init__(self, length = 1): self.length = length self.pos = np.array([0, 0]) if length > 0: self.body = Snake(length-1) self.visited = set() self.visited.add(tuple(self.pos)) def tail_visited(self): if self.length == 0: return len(self.visited) else: return self.body.tail_visited() def print(self): print(f"{self.length}: {self.pos}") #print(self.visited) if self.length > 0: self.body.print() def do(self, instruction): command = instruction.split(' ')[0] distance = int(instruction.split(' ')[1]) for _ in range(distance): self.move(command) self.body.follow(self.pos) def move(self, direction): debug_print(f"Moving {direction}") if direction == 'R': self.pos[0] += 1 elif direction == 'L': self.pos[0] -= 1 elif direction == 'U': self.pos[1] += 1 elif direction == 'D': self.pos[1] -= 1 def follow(self, new_pos): debug_print(f"Euclidian distance from T:{self.pos} to H:{new_pos} to is {np.linalg.norm(self.pos - new_pos)}") if np.linalg.norm(self.pos - new_pos) >= 2: diff = new_pos - self.pos if diff[0] > 0: self.pos[0] += 1 elif diff[0] < 0: self.pos[0] -= 1 if diff[1] > 0: self.pos[1] += 1 elif diff[1] < 0: self.pos[1] -= 1 debug_print(f"Tail moving from {self.pos}") self.visited.add(tuple(self.pos)) if self.length > 0: self.body.follow(self.pos) else: debug_print(f"Tail not moving from {self.pos}") def read_lines(filename): lines = [] with open(filename) as infile: for raw_line in infile: line = raw_line.rstrip() lines.append(line) return lines def parse_arguments(): parser = argparse.ArgumentParser(description="Script that solves the case",epilog="Have a nice day!") parser.add_argument('filename', nargs='?', default="example.txt", help='Input file') args = parser.parse_args() return args def main(): args = parse_arguments() lines = read_lines(args.filename) solve_task(lines) if __name__ == "__main__": main()