day15.py

#!/usr/bin/env python3
from __future__ import annotations

import dataclasses
import itertools
import math
import os
import time
from functools import cmp_to_key
from typing import Callable

import aocd
import parse

AOC_DAY = 15
test_data = """\
Sensor at x=2, y=18: closest beacon is at x=-2, y=15
Sensor at x=9, y=16: closest beacon is at x=10, y=16
Sensor at x=13, y=2: closest beacon is at x=15, y=3
Sensor at x=12, y=14: closest beacon is at x=10, y=16
Sensor at x=10, y=20: closest beacon is at x=10, y=16
Sensor at x=14, y=17: closest beacon is at x=10, y=16
Sensor at x=8, y=7: closest beacon is at x=2, y=10
Sensor at x=2, y=0: closest beacon is at x=2, y=10
Sensor at x=0, y=11: closest beacon is at x=2, y=10
Sensor at x=20, y=14: closest beacon is at x=25, y=17
Sensor at x=17, y=20: closest beacon is at x=21, y=22
Sensor at x=16, y=7: closest beacon is at x=15, y=3
Sensor at x=14, y=3: closest beacon is at x=15, y=3
Sensor at x=20, y=1: closest beacon is at x=15, y=3
"""


def data_parse(dstr):
    return dstr.splitlines()


test_data = data_parse(test_data)

assert len(test_data) == 14

test_answer_one = 26
test_answer_two = 56000011

input_string = "Sensor at x={sx}, y={sy}: closest beacon is at x={bx}, y={by}"
input_parser = parse.compile(input_string)


def part_one(data: list[str], target_row) -> int:

    sensors = set()
    beacons = set()
    for line in data:
        result = input_parser.parse(line)
        sx, sy, bx, by = (int(d) for d in result.named.values())
        sensors.add((sx, sy, (abs(sy - by) + abs(sx - bx))))
        beacons.add((bx, by))

    # now find non_acceptable points for beacons:
    non_acceptable = set()

    for x, y, r in sensors:
        if target_row not in range(y - r, y + r):
            continue

        distance = y - target_row
        l = r - abs(distance)
        for _x in range(x - l, x + l + 1):
            non_acceptable.add((_x, target_row))

    return len(non_acceptable - beacons)


def part_two(data: list[str], max_size) -> int:
    sensors = set()
    beacons = set()
    for line in data:
        result = input_parser.parse(line)
        sx, sy, bx, by = (int(d) for d in result.named.values())
        sensors.add((sx, sy, (abs(sy - by) + abs(sx - bx))))
        beacons.add((bx, by))

    # now find non_acceptable points for beacons:

    def excluded(x, y):
        if x > max_size or y > max_size:
            return True

        if x < 0 or y < 0:
            return True

        for xx, yy, rr in sensors:
            if (abs(xx - x) + abs(yy - y)) <= rr:
                return True

        if (x, y) in beacons:
            return True

        return False

    possible_values = set()
    for x, y, r in sensors:
        r += 1
        for i in range(r + 1):
            dx = i
            dy = r - i

            for _x, _y in [
                (x + dx, y + dy),
                (x + dx, y - dy),
                (x - dx, y + dy),
                (x - dx, y - dy),
            ]:
                if not excluded(_x, _y):
                    p = (_x, _y)
                    possible_values.add(p)

    sensors = {(x, y) for x, y, _ in sensors}
    for y in range(21):
        for x in range(21):
            if (x, y) in possible_values:
                print("#", end="")
            elif (x, y) in sensors:
                print("S", end="")
            else:
                print(" ", end="")
        print()

    assert len(possible_values) == 1, possible_values
    x, y = possible_values.pop()
    return x * 4000000 + y
    # return _x * 4000000 + _y


if __name__ == "__main__":

    part_one_ans = part_one(test_data, 10)
    assert part_one_ans == test_answer_one, f"{part_one_ans=}, not {test_answer_one=}"

    real_data = data_parse(aocd.get_data(day=AOC_DAY, year=2022))
    print("part 1:", part_one(real_data, 2000000))

    part_two_ans = part_two(test_data, max_size=20)
    assert part_two_ans == test_answer_two, f"{part_two_ans=}, not {test_answer_two=}"
    print("part 2:", part_two(real_data, max_size=4000000))