clean up code

day_01/program.py

@@ -7,31 +7,40 @@ def get_lines(filename: str) -> list:
         return [line.strip() for line in file.readlines()]
 
 
-def main():
-    # lines = get_lines("sample-input.txt")
-    lines = get_lines("input.txt")
+def get_lists(lines: list) -> (list, list):
     list_a = []
     list_b = []
     for line in lines:
-        (a, b) = line.split("   ")
+        (a, b) = line.split()
         list_a.append(int(a))
         list_b.append(int(b))
+    return list_a, list_b
+
+
+def get_total_distance(list_a: list, list_b: list) -> int:
+    total_distance = 0
     sorted_a = sorted(list_a)
     sorted_b = sorted(list_b)
-    # print(sorted_a)
-    # print(sorted_b)
+    for a, b in zip(sorted_a, sorted_b):
+        total_distance += abs(b - a)
+    return total_distance
 
-    total_distance = 0
-    for i in range(len(sorted_a)):
-        total_distance += abs(sorted_b[i] - sorted_a[i])
-    print("Part 1: Total distance is", total_distance)
 
+def get_similarity_score(list_a: list, list_b: list) -> int:
     similarity_score = 0
-    for i in range(len(list_a)):
-        search_val = list_a[i]
-        appearances = [j for j, x in enumerate(list_b) if x == search_val]
+    for search_val in list_a:
+        # returns list of positions where search_val was found
+        appearances = [i for i, x in enumerate(list_b) if x == search_val]
         similarity_score += search_val * len(appearances)
-    print("Part 2: Similarity score is", similarity_score)
+    return similarity_score
+
+
+def main():
+    # lines = get_lines("sample-input.txt")
+    lines = get_lines("input.txt")
+    list_a, list_b = get_lists(lines)
+    print("Part 1: Total distance is", get_total_distance(list_a, list_b))
+    print("Part 2: Similarity score is", get_similarity_score(list_a, list_b))
 
 
 if __name__ == '__main__':