aoc2024/day13.py

#!/usr/bin/env python3

from dataclasses import dataclass
from typing import Tuple
import re

from ortools.linear_solver import pywraplp
from ortools.math_opt.python import mathopt

@dataclass
class Claw:
   button_a: Tuple[int, int]
   button_b: Tuple[int, int]

   prize: Tuple[int, int]

claws = []
with open("day13") as f:
    l = f.readline()
    while l:
        if l == '\n':
            l = f.readline()
            continue

        button_a = re.match(r"Button A: X\+([0-9]+), Y\+([0-9]+).*", l).groups()
        button_b = re.match(r"Button B: X\+([0-9]+), Y\+([0-9]+).*", f.readline()).groups()
        prize = re.match(r"Prize: X=([0-9]+), Y=([0-9]+).*", f.readline()).groups()

        c = Claw(
            button_a=(int(button_a[0]), int(button_a[1])),
            button_b=(int(button_b[0]), int(button_b[1])),
            prize=(int(prize[0]), int(prize[1]))
        )

        claws.append(c)

        l = f.readline()


def part1():
   # (k * b_a.x) + (j * b_b.x) - x = (k * b_a.y) + (j * b_b.y) - y
   # cost = k*3 + j

   def gensol(claw: Claw):
      for k in range(100):
         for j in range(100):
            left = (k*claw.button_a[0]) + (j*claw.button_b[0]) - claw.prize[0]
            right = (k*claw.button_a[1]) + (j*claw.button_b[1]) - claw.prize[1]
            if left == 0 and right == 0:
               yield (k,j)

   def cost(k,j):
      return k*3 + j

   res = 0
   for claw in claws:
      min_cost = 500*3
      for (k,j) in gensol(claw):
         sol_cost = cost(k,j)
         print(f"Found cost for claw {claw}. k={k}, j={j}, cost={sol_cost}")
         if sol_cost < min_cost:
            min_cost = sol_cost

      if min_cost != 500*3:
         res += min_cost

   print(res)


def part2():
   # UPDATE: In fact all that below works perfectly fine if one understands
   # `add` as normal addition instead of concatenation. But for historical
   # reference:
   #
   # lots of math transformations later...
   #
   # I think this is correct. It works with part1 at least.
   #
   # It does not work with part2. I do think it could be "only" numerical errors
   # due to floating point math. Or this whole math is BS. That's very possible
   # too.
   #
   # - tried high precision math libs, didn't work
   # - tried searching around the proximity of the "exact" solution, didn't work

   # !!!!! THAT'S NOT WHAT IS MEANT WITH ADD !!!!!!!
   #
   # uncomment this to make it fail
   # for claw in claws:
   #    prize_x = int("10000000000000" + str(claw.prize[0]))
   #    prize_y = int("10000000000000" + str(claw.prize[1]))
   #    claw.prize = (prize_x, prize_y)

   for claw in claws:
      prize_x = 10000000000000 + claw.prize[0]
      prize_y = 10000000000000 + claw.prize[1]
      claw.prize = (prize_x, prize_y)

   def gensol(claw: Claw):
      a_x = claw.button_a[0]
      a_y = claw.button_a[1]

      b_x = claw.button_b[0]
      b_y = claw.button_b[1]

      p_x = claw.prize[0]
      p_y = claw.prize[1]

      epsilon = 0.00001

      y = ((a_y*p_x)-(a_x*p_y))/((b_x*a_y)-(a_x*b_y))
      print(f"{y:.200f}")

      x = ((p_x - (y*b_x))/a_x)
      print(f"{x:.200f}")

      if x>0 and abs(y - round(y)) <= epsilon and y>0 and abs(x - round(x)) <= epsilon:
         return (x, y)
      # print(f"y = (({a_y}*{p_x})-({a_x}*{p_y})) / (({b_x}*{a_y})-({a_x}*{b_y})), x = (({p_x} - ({y}*{b_x})) / {a_x})")

      # we found a solution in R but quite likely contains numerical errors
      # search the proximity for an exact match
      # epsilon = 0.0001
      # if x > 0 and y > 0:
      #    for i in range(-100,100):
      #       for j in range(-100,100):
      #          if abs(((x+i) * a_x) + ((y+j) * b_x) - p_x) < epsilon and abs(((x+i) * a_y) + ((y+j) * b_y) - p_y) < epsilon:
      #             return (int(x),int(y))

   def cost(k,j):
      return k*3 + j

   res = 0
   for i,claw in enumerate(claws):
      print(f"running claw {i+1} of {len(claws)}")
      sol = gensol(claw)
      if sol:
         k,j = sol
         sol_cost = cost(k,j)
         print(f"Found cost for claw {claw}. k={k}, j={j}, cost={sol_cost}")
         print(f"{claw.prize[0]}={k*claw.button_a[0]+j*claw.button_b[0]}")
         res += sol_cost

   print(res)

def part2_get_the_bfg():
   sol = 0
   for i, claw in enumerate(claws):
      solver = pywraplp.Solver.CreateSolver("SAT")
      if not solver:
         return

      infinity = solver.infinity()
      x = solver.IntVar(lb=0.0, ub=infinity, name="x")
      y = solver.IntVar(lb=0.0, ub=infinity, name="y")
      solver.Add(x*claw.button_a[0] + y*claw.button_b[0] <= claw.prize[0])
      solver.Add(x*claw.button_a[0] + y*claw.button_b[0] >= claw.prize[0])
      solver.Add(x*claw.button_a[1] + y*claw.button_b[1] <= claw.prize[1])
      solver.Add(x*claw.button_a[1] + y*claw.button_b[1] >= claw.prize[1])
      solver.Minimize(3*x + y)

      status = solver.Solve()
      if status == pywraplp.Solver.OPTIMAL:
         x = x.solution_value()
         y = y.solution_value()
         sol += solver.Objective().Value()

         print(f"Claw {i: >3}: {x=:>20.3f}\t{y=:>20.3f}\t{solver.Objective().Value()=:.2f} [{claw}]")
      else:
         print(f"Claw {i: >3}: no solution")

   print(f"{sol:.5f}")

def part2_get_the_bfg2():
   sol = 0
   for i, claw in enumerate(claws):
      model = mathopt.Model()
      x = model.add_variable(lb=0.0, name="x")
      y = model.add_variable(lb=0.0, name="y")

      model.add_linear_constraint(x*claw.button_a[0] + y*claw.button_b[0] <= claw.prize[0])
      model.add_linear_constraint(x*claw.button_a[0] + y*claw.button_b[0] >= claw.prize[0])
      model.add_linear_constraint(x*claw.button_a[1] + y*claw.button_b[1] <= claw.prize[1])
      model.add_linear_constraint(x*claw.button_a[1] + y*claw.button_b[1] >= claw.prize[1])

      model.minimize(3*x + y)
      result = mathopt.solve(model, mathopt.SolverType.GLOP)
      if result.termination.reason != mathopt.TerminationReason.OPTIMAL:
         print(f"model failed to solve: {result.termination}")

      else:
         print("MathOpt solve succeeded")
         print("Objective value:", result.objective_value())
         print("x:", result.variable_values()[x])
         print("y:", result.variable_values()[y])

         x = result.variable_values()[x]
         y = result.variable_values()[y]

         epsilon = 0.001
         if x>0 and y>0 and abs(x - round(x)) <= epsilon and abs(y - round(y)) <= epsilon:
            sol += result.objective_value()

   print(f"{sol:.5f}")