AoC 2024 Day 8

from itertools import combinations
from collections import defaultdict


with open("day8input.txt") as f:
    input_lines = f.read().splitlines()


# part 1
def find_num_antennas(grid: list[str]) -> int:
    antenna_locations = defaultdict(list)

    row_len = len(grid)
    col_len = len(grid[0])

    # find all antenna locations, group them in the dictionary by same value
    for y in range(row_len):
        for x in range(col_len):
            if input_lines[y][x] != '.':
                antenna_locations[grid[y][x]].append((y, x))

    # set is used to store found antinode locations to avoid double counting
    output = set()
    # loop through each found frequency
    for key in antenna_locations:
        # find all combinations of two antennas with this frequency
        all_combos = list(combinations(antenna_locations[key], 2))
        for combo in all_combos:
            # find the offset between the two points and the antinodes on each side
            diff = combo[1][0] - combo[0][0], combo[1][1] - combo[0][1]
            point1 = combo[0][0] - diff[0], combo[0][1] - diff[1]
            point2 = combo[1][0] + diff[0], combo[1][1] + diff[1]
            # increment output if antinode location is in bounds
            if 0 <= point1[0] < row_len and 0 <= point1[1] < col_len:
                output.add(point1)
            if 0 <= point2[0] < row_len and 0 <= point2[1] < col_len:
                output.add(point2)

    return len(output)


print(find_num_antennas(input_lines))


# part 2
def find_num_antennas2(grid: list[str]) -> int:
    antenna_locations = defaultdict(list)

    row_len = len(grid)
    col_len = len(grid[0])

    # find all antenna locations, group them in the dictionary by same value
    for y in range(row_len):
        for x in range(col_len):
            if input_lines[y][x] != '.':
                antenna_locations[grid[y][x]].append((y, x))

    # set is used to store found antinode locations to avoid double counting
    output = set()
    # loop through each found frequency
    for key in antenna_locations:
        # find all combinations of two antennas with this frequency
        all_combos = list(combinations(antenna_locations[key], 2))
        for combo in all_combos:
            # find the offset between the two points
            diff = combo[1][0] - combo[0][0], combo[1][1] - combo[0][1]
            point1 = combo[0]
            point2 = combo[1]
            # for all points that are the same offset away from one of the original two points or newly found ones,
            # count it if it is in bounds
            while 0 <= point1[0] < row_len and 0 <= point1[1] < col_len:
                output.add(point1)
                point1 = point1[0] - diff[0], point1[1] - diff[1]
            while 0 <= point2[0] < row_len and 0 <= point2[1] < col_len:
                output.add(point2)
                point2 = point2[0] + diff[0], point2[1] + diff[1]

    return len(output)


print(find_num_antennas2(input_lines))