Advent of Code - Day 1 - v1

day1.py

import argparse
import sys
import re
from collections import defaultdict
from io import StringIO


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-i",
        "--input-file",
        type=str,
        default="-",
        help="Input file, put - for stdin",
    )

    return parser.parse_args()


def parse_input(f: StringIO) -> tuple[list[int], list[int]]:
    list1 = []
    list2 = []

    for line in f:
        if m := re.match(r'^([0-9]+)\s+([0-9]+)$', line):
            list1.append(int(m.group(1)))
            list2.append(int(m.group(2)))
        else:
            raise ValueError(f"Invalid line: {line}")

    return list1, list2


def compute_distance(list1: list[int], list2: list[int]) -> int:
    s1 = sorted(list1)
    s2 = sorted(list2)

    return sum(abs(s1[i] - s2[i]) for i in range(len(list1)))


def compute_similarity(list1: list[int], list2: list[int]) -> float:
    right_count = defaultdict(int)

    for x in list2:
        right_count[x] += 1

    return sum(right_count[x] * x for x in list1)


def main():
    args = parse_args()

    if args.input_file == "-":
        input_file = sys.stdin
    else:
        input_file = open(args.input_file, "r")

    list1, list2 = parse_input(input_file)
    distance = compute_distance(list1, list2)
    similarity = compute_similarity(list1, list2)

    print(f"Distance: {distance}")
    print(f"Similarity: {similarity}")


if __name__ == "__main__":
    main()