AOC 2021 Day 3

# aoc202103.py

import pathlib
import sys
from typing import Any


def parse(puzzle_input: str) -> list[list[str]]:
    """ Parse input """
    # Return a 2-D array
    return [[c for c in line] for line in puzzle_input.splitlines()]


def transpose_array(array: list[list[Any]]) -> list[list[Any]]:
    """ Transpose a 2-D array of any type """
    return list(zip(*array))


def most_frequent(seq: list[str], preferred: str) -> str:
    """ Return most frequent item in a list seq """
    first, sec = list(set(seq))
    # Deal with equal count
    if seq.count(first) == seq.count(sec):
        return preferred
    return max(set(seq), key=seq.count)


def least_frequent(seq: list[str], preferred: str) -> str:
    """ Return most frequent item in a list seq """
    first, sec = list(set(seq))
    # Deal with equal count
    if seq.count(first) == seq.count(sec):
        return preferred
    return min(set(seq), key=seq.count)


def binary_product(a: str, b: str) -> int:
    """ Take two binary numbers as strings and return their product as an int """
    return int(a, 2) * int(b, 2)


def flatten(t: list[list[Any]]) -> list[Any]:
    """ Flatten a nested list """
    return [item for sublist in t for item in sublist]


def get_rating(array: list[list[str]],
               columns: int,
               low: bool = False) -> list[str]:
    """ Find one row in an array that matches the bit_criteria """
    for col in range(columns):
        cols = transpose_array(array)
        # For the CO2 scrubber rating
        if low:
            bit = least_frequent(cols[col], "0")
        # For the oxygen generator rating
        else:
            bit = most_frequent(cols[col], "1")
        # Only include the rows which meet the criteria
        target = [row for row in array if row[col] == bit]
        array = target[:]
        # Bail out early if target identified
        if len(array) == 1:
            break
    # Just in case...
    if len(array) != 1:
        raise ValueError("Unique target not found")
    # Unnest the list and return
    return flatten(target)


def part1(data: list[list[str]]) -> int:
    """ Solve part 1 """
    # Transpose the array
    cols = transpose_array(data)
    # Get the most frequent values from cols
    gamma = ''.join(most_frequent(col, '1') for col in cols)
    # Get the least frequent values from cols
    epsilon = ''.join(least_frequent(col, '0') for col in cols)
    return binary_product(gamma, epsilon)


def part2(data: list[str]) -> int:
    """ Solve part 2 """
    # Transpose the array
    cols = transpose_array(data)
    # Get the oxygen generator rating
    oxy = ''.join(get_rating(data, len(cols)))
    # Get the CO2 scrubber rating
    co2 = ''.join(get_rating(data, len(cols), low=True))
    return binary_product(oxy, co2)


def solve(puzzle_input: str) -> tuple[int, int]:
    """ Solve the puzzle for the given input """
    data = parse(puzzle_input)
    solution1 = part1(data)  # Correct answer was 2498354 (with my data)
    solution2 = part2(data)  # Correct answer was 3277956 (with my data)

    return solution1, solution2


if __name__ == "__main__":
    for path in sys.argv[1:]:
        print(f"{path}:")
        puzzle_input = pathlib.Path(path).read_text().strip()
        solutions = solve(puzzle_input)
        print('\n'.join(str(solution) for solution in solutions))