Spaces:
Running
Running
File size: 1,709 Bytes
a4da721 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import numpy as np
def parse_block(block):
# Parse button A
A = block.split("\n")[0]
_, ax, ay = A.split("+")
ax = ax.split(",")[0]
# Parse button B
B = block.split("\n")[1]
_, bx, by = B.split("+")
bx = bx.split(",")[0]
# Parse prize loc
P = block.split("\n")[2]
_, px, py = P.split("=")
px = px.split(",")[0]
ax, ay = int(ax), int(ay)
bx, by = int(bx), int(by)
px, py = int(px), int(py)
return np.array([[ax, bx], [ay, by]]), np.array([[px], [py]])
def load_data(file):
with open(file) as f:
data = f.read()
blocks = [parse_block(block) for block in data.split("\n\n")]
return blocks
def solve_equation(W, Y):
# Each block consists of W, Y matrices, we solve the question
# X = W^-1 * Y
X = (np.linalg.inv(W) @ Y)
return X
def is_valid(a):
return np.isclose(a, round(a)) and a <= 100
def solve(file):
blocks = load_data(file)
total = 0
for block in blocks:
W, Y = block
X = solve_equation(W, Y)
a, b = X[0][0], X[1][0]
if is_valid(a) and is_valid(b):
total += round(a)* 3 + round(b)*1
print((total))
## Part 1
file = "input.txt"
solve(file)
## Part 2
def is_valid(a):
return np.isclose(a, round(a), atol=0.01, rtol=0) and a >= 0
def solve_2(file, offset):
blocks = load_data(file)
total = 0
for block in blocks:
W, Y = block
Y = Y + offset
X = solve_equation(W, Y)
a, b = X[0][0], X[1][0]
if is_valid(a) and is_valid(b):
total += round(a)* 3 + round(b)* 1
print((total))
file = "input.txt"
offset = 10000000000000
solve_2(file, offset)
|