day10b.py

import numpy as np
import cvxpy as cp

f = open('input.txt', 'r')
lines = [line.split(' ') for line in f.read().split('\n')]
button_lists = [[list(map(int, btn[1:-1].split(','))) for btn in line[1:-1]] for line in lines]
joltages = [list(map(int, line[-1][1:-1].split(','))) for line in lines]


res = 0
for buttons, joltage in zip(button_lists, joltages):
	m = len(joltage)
	n = len(buttons)

	A = np.zeros((m, n))
	for col, btn in enumerate(buttons):
		for row in btn:
			A[row][col] = 1

	b = np.array(joltage)

	x = cp.Variable(n, integer=True, nonneg=True)

	constraints = [A @ x == b]
	objective = cp.Minimize(cp.norm1(x))
	prob = cp.Problem(objective, constraints)

	prob.solve()

	res += int(sum(x.value))

print(res)