Advent of Code - Day 1 - v3

day1.py

import argparse
import sys
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:
        num1, num2 = parse_line(line.strip())
        if num1 is not None and num2 is not None:
            list1.append(num1)
            list2.append(num2)
        else:
            raise ValueError(f"Invalid line: {line}")

    return list1, list2


def parse_line(line: str) -> tuple[int | None, int | None]:
    state = 0
    num1 = num2 = ''

    for char in line:
        if state == 0:
            if char.isdigit():
                num1 += char
                state = 1
            else:
                return None, None
        elif state == 1:
            if char.isdigit():
                num1 += char
            elif char.isspace():
                state = 2
            else:
                return None, None
        elif state == 2:
            if char.isdigit():
                num2 += char
                state = 3
            elif char.isspace():
                pass
            else:
                return None, None
        elif state == 3:
            if char.isdigit():
                num2 += char
            else:
                return None, None

    if state == 3 and num1 and num2:
        return int(num1), int(num2)
    else:
        return None, None


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()