adventofcode2023/day03-2.py

from helpers import debug, load_input

def is_digit(n):
    try:
        m = int(n)
    except ValueError:
        return False
    return True

def is_star(n):
    return n == "*"

input_numbers = []
stars = []

class InputNumber:
    def __init__(self, number, start_pos, end_pos, line):
        self.number = number
        self.start_pos = start_pos
        self.end_pos = end_pos
        self.line = line

    def _is_symbol(self, char):
        return (char != "." and not is_digit(char))

    def _check_up(self, input):
        if self.line == 0:
            return False
        cur_line = input[self.line - 1]
        for i in range(max(0, self.start_pos-1), min(len(cur_line)-1, self.end_pos+2)):
            char = cur_line[i]
            if self._is_symbol(char):
                debug(f"Found symbol next to {self.number} at line {self.line-1}: {i}")
                return True
        return False
    
    def _check_down(self, input):
        if self.line == len(input)-1:
            return False
        cur_line = input[self.line + 1]
        for i in range(max(0, self.start_pos-1), min(len(cur_line)-1, self.end_pos+2)):
            char = cur_line[i]
            if self._is_symbol(char):
                debug(f"Found symbol next to {self.number} at line {self.line+1}: {i}")
                return True
        return False
    
    def _check_left(self, input):
        if self.start_pos == 0:
            debug("Already at the left-hand side of the line")
            return False
        if self._is_symbol(input[self.line][self.start_pos-1]):
            debug(f"Found symbol next to {self.number} at line {self.line}: {self.start_pos-1}")
            return True
        return False

    def _check_right(self, input):
        if self.end_pos == len(input[0]) - 1:
            debug("Already at the right-hand side of the line")
            return False
        if self._is_symbol(input[self.line][self.end_pos+1]):
            debug(f"Found symbol next to {self.number} at line {self.line}: {self.end_pos+1}")
            return True
        return False
    
    def is_part_number(self, input):
        return (
            self._check_up(input)
            or self._check_down(input)
            or self._check_left(input)
            or self._check_right(input)
        )
    
    def thats_me(self, line, pos):
        return self.line == line and pos in range(self.start_pos, self.end_pos+1)

class Star:
    def __init__(self, line, pos):
        self.line = line
        self.pos = pos
        self.adjacent_numbers = []

    def _check_up(self, input):
        if self.line == 0:
            debug("Already at the top")
            return False
        found_one = False
        for i in range(self.pos-1, self.pos+2):
            debug(f"Checking {self.line-1}:{i}")
            char = input[self.line-1][i]
            if is_digit(char):
                target_number = [number for number in input_numbers if number.thats_me(self.line-1, i)][0]
                if target_number not in self.adjacent_numbers:
                    self.adjacent_numbers.append(target_number)
                found_one = True
        return found_one
    
    def _check_down(self, input):
        if self.line == len(input) - 1:
            debug("Already at the bottom")
            return False
        found_one = False
        for i in range(self.pos-1, self.pos+2):
            debug(f"Checking {self.line+1}:{i}")
            char = input[self.line+1][i]
            if is_digit(char):
                target_number = [number for number in input_numbers if number.thats_me(self.line+1, i)][0]
                if target_number not in self.adjacent_numbers:
                    self.adjacent_numbers.append(target_number)
                found_one = True
        return found_one
    
    def _check_left(self, input):
        if self.pos == 0:
            debug ("Already at the left side")
            return False
        debug(f"Checking {self.line}:{self.pos-1}")
        char = input[self.line][self.pos-1]
        if is_digit(char):
            target_number = [number for number in input_numbers if number.thats_me(self.line, self.pos-1)][0]
            if target_number not in self.adjacent_numbers:
                self.adjacent_numbers.append(target_number)
            return True
        return False
    
    def _check_right(self, input):
        if self.pos == len(input[0]) - 1:
            debug ("Already at the right side")
            return False
        debug(f"Checking {self.line}:{self.pos+1}")
        char = input[self.line][self.pos+1]
        if is_digit(char):
            target_number = [number for number in input_numbers if number.thats_me(self.line, self.pos+1)][0]
            if target_number not in self.adjacent_numbers:
                self.adjacent_numbers.append(target_number)
            return True
        return False
    
    def _is_gear(self, input):
        self._check_up(input)
        self._check_down(input)
        self._check_left(input)
        self._check_right(input)
        l_adj = len(self.adjacent_numbers)
        n_adj = [number.number for number in self.adjacent_numbers]
        debug(f"Star at {self.line}:{self.pos} has {l_adj} adjacent number{'s' if l_adj > 1 else ''}: {n_adj}")
        return len(self.adjacent_numbers) == 2
    
    def gear_value(self, input):
        if self._is_gear(input):
            return self.adjacent_numbers[0].number * self.adjacent_numbers[1].number
        return 0

def main():
    lines = load_input(3)
    # Input is lines of periods (.), numbers (0-9), and symbols 
    # (anything that isn't a period or number)
    # Any number adjacent (horizontally, vertically, diagonally)
    # to a symbol is a part number. Get the sum of the part numbers.
    max_len = len(lines[0]) # all lines are the same length
    for j, line in enumerate(lines):
        i = 0
        while i < max_len:
            if is_star(line[i]):
                star = Star(
                    line=j,
                    pos=i
                )
                stars.append(star)
                # debug(f"Found a star at {star.line}:{star.pos}")
            elif is_digit(line[i]):
                current_number = ""
                start_pos = i
                while i < max_len and is_digit(line[i]):
                    current_number += f"{line[i]}"
                    i += 1
                if i < max_len and is_star(line[i]):
                    star = Star(
                        line=j,
                        pos=i
                    )
                    stars.append(star)
                    # debug(f"Found a star next to a number at {star.line}:{star.pos}")
                end_pos = i
                input_number = InputNumber(
                    number = int(current_number),
                    start_pos = start_pos,
                    end_pos = end_pos-1, # i is one greater than the end of the number
                    line = j
                )
                input_numbers.append(input_number)
                # debug(f"Found {input_number.number} at line {input_number.line} ({input_number.start_pos}, {input_number.end_pos})")
            i += 1
    # Now I have a list of the positions of all the numbers and stars in the input.
    gear_values = []
    for s in stars:
        gear_values.append(s.gear_value(lines))
    
    print(sum(gear_values))

if __name__ == "__main__":
    main()