aoc_2023_day_3_part_1.py

from dataclasses import dataclass
from typing import Dict, Set


@dataclass
class PartNumber:
    id: int
    value: int

    def __hash__(self):
        return self.id


_part_number_id = 1


def get_part_number_id() -> int:
    """Global unique counter for part number id"""
    global _part_number_id
    _part_number_id += 1
    return _part_number_id - 1


def parse_numbers_per_line(raw_line, y_index) -> Dict[str, PartNumber]:
    """
    Map the numbers to the lookup table where
    key = x_index, y_index
    value = PartNumber object

    multiple keys can map to the same PartNumber object
    """
    pointer = 0
    partial_lookup_table = {}

    def create_key(x_index):
        return f'{x_index},{y_index}'

    while pointer < len(raw_line):

        if raw_line[pointer].isdigit():
            # start a new key for each x,y coordinate
            keys = set()

            start_pointer = pointer
            keys.add(create_key(pointer))

            while (
                    pointer < len(raw_line) and
                    raw_line[pointer].isdigit()
            ):
                keys.add(create_key(pointer))
                pointer += 1

            number = PartNumber(
                id=get_part_number_id(),
                value=int(raw_line[start_pointer:pointer])
            )

            for key in keys:
                partial_lookup_table[key] = number

        else:
            pointer += 1

    return partial_lookup_table


def find_adjacent(lookup_table: dict, x :int, y: int) -> Set[PartNumber]:
    """
    Find left, right, up, down, diagonal up left,
    diagonal up right, diagonal down left, diagonal down right
    return a set of unique of PartNumber objects
    """

    adjacent = set()
    for x_offset in range(-1, 2):
        for y_offset in range(-1, 2):
            if x_offset == 0 and y_offset == 0:
                continue

            key = f'{x + x_offset},{y + y_offset}'
            if found := lookup_table.get(key):
                adjacent.add(found)

    return adjacent


def solve():
    """
    Make a lookup table of all the parts where the key is the 
    x,y and the value is the `PartNumber` instance.

    Since length coordinate of a part can span multiple x coordinates,
    The lookup table allows multiple keys to refer to the same object.

    example: 

    partA = PartNumber(id=1,value=35)

    {
        '2,2': partA 
        '3,2': partA
    }

    when lookup 2,2 and 3,2, they both refer to the same object id=1 and value=35
    When later summed up, we can use a set to remove duplicates

    total = set([partA, partA, partB]) -> {partA, partB}
    """
    lookup_table_result: Dict[str, PartNumber] = {}

    # enumerate each of part numbers to a unique id
    for y_idx, line in enumerate(_input.split('\n')):
        lookup_table_result.update(
            parse_numbers_per_line(line, y_idx)
        )

    considered_parts: Set[PartNumber] = set()

    # find all the parts that are adjacent to a non-digit
    for y_idx, line in enumerate(_input.split('\n')):
        for x_idx, char in enumerate(line):
            if char != '.' and not char.isdigit():
                considered_parts.update(
                    find_adjacent(lookup_table_result, x_idx, y_idx)
                )

    print(sum(
        part.value
        for part in considered_parts
    ))


if __name__ == '__main__':
    solve()