import os, argparse import numpy as np debug_print = print debug_print = lambda x: x def solve_task(lines): s = Snake() s.print() for line in lines: s.do(line) #s.print() print(len(s.visited)) class Snake: def __init__(self): self.H = np.array([0, 0]) self.last_H = self.H self.T = np.array([0, 0]) self.visited = set() self.visited.add((0,0)) def print(self): print(f"H: {self.H}, T: {self.T}") print(self.visited) def do(self, instruction): command = instruction.split(' ')[0] distance = int(instruction.split(' ')[1]) for _ in range(distance): self.go(command) self.adjust_tail() def go(self, direction): debug_print(f"Moving {direction}") self.last_H = self.H.copy() if direction == 'R': self.H[0] += 1 elif direction == 'L': self.H[0] -= 1 elif direction == 'U': self.H[1] += 1 elif direction == 'D': self.H[1] -= 1 def adjust_tail(self): debug_print(f"Euclidian distance from T:{self.T} to H:{self.H} is {np.linalg.norm(self.H - self.T)}") if np.linalg.norm(self.H - self.T) >= 2: debug_print(f"Tail moving from {self.T} to {self.last_H}") self.T = self.last_H self.visited.add(tuple(self.T)) else: debug_print(f"Tail not moving from {self.T} to {self.last_H}") 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()