advent-of-code/2022/python/day15.py

import matrix
import math
import sys
from pprint import pprint
import shared
from scanf import scanf
from typing import Optional, List, Tuple
from dataclasses import dataclass
from collections import defaultdict


def cityblock(y1, x1, y2, x2):
    return abs(y2 - y1) + abs(x2 - x1)


@dataclass
class Sensor:
    sX: int
    sY: int
    bX: int
    bY: int
    limit: Tuple[int, int] = (0, 0)
    _d: int = None
    _edges: List[Tuple[int, int]] = None
    _border: List[Tuple[int, int]] = None

    def __str__(self):
        return (
            f"Sensor(sX={self.sX}, sY={self.sY}, bX={self.bX},"
            f"bY={self.bY}, d={self._d}, edges={len(self._edges)}, borders={len(self._borders)})"
        )

    @property
    def s(self):
        return (self.sY, self.sX)

    @property
    def b(self):
        return (self.bY, self.bX)

    @property
    def distance(self):
        if self._d is None:
            self._d = cityblock(self.sY, self.sX, self.bY, self.bX)
        return self._d

    def distance_to(self, bY, bX):
        return cityblock(self.sY, self.sX, bY, bX)

    def on_line(self, y):
        midpoint = (y, self.s[1])
        d = self.distance_to(*midpoint)
        if d > self.distance:
            return []

        need = self.distance - d

        start = (y, midpoint[1] - need)
        end = (y, midpoint[1] + need)

        return list(range(start[1], end[1] + 1))

    def in_range(self, bY, bX):
        d = cityblock(self.sY, self.sX, bY, bX)
        if self.d < d:
            return False
        return True

    def in_diamond(self):
        sX, sY = self.sX, self.sY
        up_lim = sY - self.distance
        dn_lim = sY + self.distance

        le_lim = sX - self.distance
        ri_lim = sX + self.distance

        u = (up_lim, sX)
        d = (dn_lim, sX)
        l = (sY, le_lim)
        r = (sY, ri_lim)

        infliction = 1
        height = -1
        for idx, x in enumerate(range(l[1], r[1] + 1)):
            height += infliction
            if (sY, x) == self.s:
                infliction = -1
            for y in range(sY - height, sY + height + 1):
                yield (y, x)

    def edges(self):
        if self._edges:
            return self._edges
        sX, sY = self.sX, self.sY
        up_lim = sY - self.distance
        dn_lim = sY + self.distance

        le_lim = sX - self.distance
        ri_lim = sX + self.distance

        u = (up_lim, sX)
        d = (dn_lim, sX)
        l = (sY, le_lim)
        r = (sY, ri_lim)

        infliction = 1
        height = -1
        edges = set()
        # to left -1  and right + 1
        for idx, x in enumerate(range(l[1], r[1] + 1)):
            height += infliction
            if (sY, x) == self.s:
                infliction = -1
            edges.add((sY - height, x))
            edges.add((sY + height, x))
        self._edges = edges
        return self._edges

    def border(self):
        if self._border:
            return self._border
        sX, sY = self.sX, self.sY
        up_lim = sY - self.distance
        dn_lim = sY + self.distance

        le_lim = sX - self.distance
        ri_lim = sX + self.distance

        u = (up_lim, sX)
        d = (dn_lim, sX)
        l = (sY, le_lim)
        r = (sY, ri_lim)

        infliction = 1
        height = -1
        border = set()
        # to left -1  and right + 1
        for idx, x in enumerate(range(l[1] - 1, r[1] + 2)):
            height += infliction
            if (sY, x) == self.s:
                infliction = -1
            border.add((sY - height, x))
            border.add((sY + height, x))
        self._border = border
        return self._border


def part1(rows, sample=False):
    sensors = []
    sensor_points = []
    beacon_points = []
    ineligible_points = set()
    xSet = set()
    ySet = set()
    for row in rows:
        x, y, bx, by = scanf(
            "Sensor at x=%d, y=%d: closest beacon is at x=%d, y=%d", row
        )
        xSet.add(x)
        xSet.add(bx)
        ySet.add(y)
        ySet.add(by)
        sensors.append(Sensor(sX=x, sY=y, bX=bx, bY=by, limit=(0, 0)))
    minX, maxX = min(xSet), max(xSet)
    minY, maxY = min(ySet), max(ySet)
    limLo = min(minX, minY)
    limHi = max(maxX, maxY)

    for sensor in sensors:
        sensor.limit = (limLo, limHi)
        sensor_points.append(sensor.s)
        beacon_points.append(sensor.b)
        if sample:
            for yx in sensor.in_diamond():
                ineligible_points.add(yx)

    CHECK_ROW = 2000000
    if sample:
        CHECK_ROW = 10

    ineligible = set()
    for s in sensors:
        coll = s.on_line(CHECK_ROW)
        ineligible.update(coll)

    count_ignoring_current_beacons = 0
    for i in ineligible:
        if (CHECK_ROW, i) not in beacon_points:
            count_ignoring_current_beacons += 1

    print(count_ignoring_current_beacons, "with removing beacons, final answer")
    if not sample:
        return

    mx = matrix.matrix_of_size(maxX + 1, maxY + 1)
    for yx in ineligible_points:
        y, x = yx
        if y >= 0 and x >= 0:
            if y <= maxY and x <= maxX:
                mx[y][x] = "#"
    for yx in beacon_points:
        y, x = yx
        if y >= 0 and x >= 0:
            if y <= maxY and x <= maxX:
                mx[y][x] = "B"
    for yx in sensor_points:
        y, x = yx
        if y >= 0 and x >= 0:
            if y <= maxY and x <= maxX:
                mx[y][x] = "S"
    print(matrix.ppmx(mx, pad=False, space=True))


tuning = lambda y, x: y + (4000000 * x)


def part2(rows, sample=False):
    sensors = []
    sensor_points = []
    beacon_points = []
    ineligible_points = set()
    xSet = set()
    ySet = set()
    for row in rows:
        x, y, bx, by = scanf(
            "Sensor at x=%d, y=%d: closest beacon is at x=%d, y=%d", row
        )
        xSet.add(x)
        xSet.add(bx)
        ySet.add(y)
        ySet.add(by)
        sensors.append(Sensor(sX=x, sY=y, bX=bx, bY=by))
    minX, maxX = min(xSet), max(xSet)
    minY, maxY = min(ySet), max(ySet)

    for sensor in sensors:
        _ = sensor.edges()
        _ = sensor.border()
        sensor_points.append(sensor.s)
        beacon_points.append(sensor.b)
        if sample:
            for yx in sensor.in_diamond():
                ineligible_points.add(yx)

    L = 4000000
    if sample:
        L = 20

    borders = defaultdict(int)
    for s in sensors:
        for yx in s.border():
            y, x = yx
            if y > 0 and y <= L and x > 0 and x <= L:
                borders[yx] += 1

    TARGET = None
    for (eY, eX) in borders.keys():
        # print("checking:",(eY,ex))
        away_from = []
        for idx, s in enumerate(sensors):
            d = s.distance_to(eY, eX)
            if d > s.distance:
                away_from.append(s.s)
        if len(away_from) == len(sensors):
            TARGET = (eY, eX)
            print(TARGET, tuning(eY, eX))
            break

    if not sample:
        return

    # """ PRINT OUTPUT """

    mx = matrix.matrix_of_size(maxX + 1, maxY + 1)
    for yx in ineligible_points:
        y, x = yx
        if y >= 0 and x >= 0:
            if y <= maxY and x <= maxX:
                mx[y][x] = "#"

    for yx in beacon_points:
        y, x = yx
        if y >= 0 and x >= 0:
            if y <= maxY and x <= maxX:
                mx[y][x] = "B"

    for yx in sensor_points:
        y, x = yx
        if y >= 0 and x >= 0:
            if y <= maxY and x <= maxX:
                mx[y][x] = "S"

    mx[TARGET[0]][TARGET[1]] = "!"
    matrix.highlight(
        mx,
        blink_green=[
            TARGET,
        ],
    )


def main():
    sample = False
    if sys.argv[-1] == "--sample":
        sample = True
    rows = [row for row in shared.load_rows(15)]

    with shared.elapsed_timer() as elapsed:
        part1(rows, sample)
        print("🕒", elapsed())

    with shared.elapsed_timer() as elapsed:
        part2(rows, sample)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()