Optimized

1 files changed, 11 insertions, 16 deletions

diff --git a/day5/__init__.py b/day5/__init__.py
index 030df19..0e48b16 100644
--- a/day5/__init__.py
+++ b/day5/__init__.py
@@ -6,6 +6,8 @@ from aoc import BaseAssignment
 Vector = Tuple[Tuple[int, int], Tuple[int, int]]
 
 class Assignment(BaseAssignment):
+    including_diagonals = False
+
     def parse_item(self, item: str) -> Vector:
         start, end = item.split('->')
         return (
@@ -16,6 +18,14 @@ class Assignment(BaseAssignment):
     def read_input(self, example = False) -> List[Vector]:
         return list(super().read_input(example))
 
+    def run(self, input: List[Vector]) -> int:
+        coordinates = []
+
+        for vector in input:
+            coordinates += points_in_vector(vector, self.including_diagonals)
+
+        return len([c for c in Counter(coordinates).values() if c >= 2])
+
 
 def is_horizontal(vector: Vector):
     return vector[0][0] == vector[1][0]
@@ -70,22 +80,7 @@ def points_in_vector(vector: Vector, includes_diagonals: bool = False) -> List[T
 class AssignmentOne(Assignment):
     example_result = 5
 
-    def run(self, input: List[Vector]) -> int:
-        coordinates = []
-
-        for vector in input:
-            coordinates += points_in_vector(vector)
-
-        return len([c for c in Counter(coordinates).values() if c >= 2])
-
 
 class AssignmentTwo(Assignment):
     example_result = 12
-
-    def run(self, input: List[Vector]) -> int:
-        coordinates = []
-
-        for vector in input:
-            coordinates += points_in_vector(vector, True)
-
-        return len([c for c in Counter(coordinates).values() if c >= 2])
\ No newline at end of file
+    including_diagonals = True