aoc

My solutions for the Advent of Code
git clone https://git.tronto.net/aoc
Download | Log | Files | Refs | README

b.py (3650B)


      1 import fileinput
      2 from math import gcd
      3 
      4 inf = 999999999
      5 
      6 # Read an input line and return a triple A, b, c where:
      7 #   A is the matrix of coefficient (A[i][j] is 1 if pressing the j-th button
      8 #     increases the i-th counter, 0 otherwise).
      9 #   b is the list of required final values of each counter.
     10 #   c is a list of bounds for the number of times each button can be pressed:
     11 #     c[i] is the minimum required value of a counter increased by button i.
     12 def readline(l):
     13 	b = [int(x) for x in l[l.index('{')+1:l.index('}')].split(',')]
     14 
     15 	A = [[] for j in b]
     16 	c = []
     17 
     18 	end = l.index(']')
     19 	while '(' in l[end:]:
     20 		begin = l[end:].index('(') + end + 1
     21 		end = l[begin:].index(')') + begin
     22 		bu = [int(x) for x in l[begin:end].split(',')]
     23 		c.append(min(b[j] for j in bu))
     24 		for i in range(len(b)):
     25 			A[i].append(1 if i in bu else 0)
     26 
     27 	return A, b, c
     28 
     29 # swaprow, swapcol and reducerow are utility functions for the matrix
     30 # row reduction procedure.
     31 def swaprow(A, b, i, j):
     32 	if i != j:
     33 		A[i], A[j] = A[j], A[i]
     34 		b[i], b[j] = b[j], b[i]
     35 
     36 def swapcol(A, c, i, j):
     37 	if i != j:
     38 		for k in range(len(A)):
     39 			A[k][i], A[k][j] = A[k][j], A[k][i]
     40 		c[i], c[j] = c[j], c[i]
     41 
     42 def reducerow(A, b, i, j):
     43 	if A[i][i] != 0:
     44 		x = A[i][i]
     45 		y = -A[j][i]
     46 		d = gcd(x, y)
     47 		A[j] = [(y*A[i][k]+x*A[j][k])//d for k in range(len(A[i]))]
     48 		b[j] = (y*b[i]+x*b[j])//d
     49 
     50 # Reduce the matrix A by row, including back subsitution. Returns a triple
     51 # A, b, c where:
     52 #   A is the reduced matrix in the form [D|p] with D diagonal and p a matrix
     53 #     whose number of columns is the number of free parameters for Ax = b.
     54 #   b and c are the adjusted versions of the input paramters with the same
     55 #     names. In particular, b is adjusted every time a row is reduced and
     56 #     when rows are swapped, while c (the list of bounds) is adjusted when
     57 #     two columns are swapped (which is equivalent to changing the order of
     58 #     two buttons).
     59 def reduce(A, b, c):
     60 	for i in range(len(A[0])):
     61 		# Swap columns until there is one in position i with at least
     62 		# one non-zero element.
     63 		I = []
     64 		k = i
     65 		while len(I) == 0 and k < len(A[0]):
     66 			swapcol(A, c, i, k)
     67 			I = [j for j in range(i, len(A)) if A[j][i] != 0]
     68 			k += 1
     69 
     70 		# If no such column is found, we are done
     71 		if len(I) == 0:
     72 			break
     73 
     74 		# Swap rows so that A[i][i] is non-zero
     75 		swaprow(A, b, i, I[0])
     76 
     77 		# Reduce all other rows
     78 		for j in range(i+1, len(A)):
     79 			reducerow(A, b, i, j)
     80 
     81 	# Remove all rows of zero and check if the system is solvable
     82 	I = [i for i in range(len(A)) if any(a != 0 for a in A[i])]
     83 	A = [A[i] for i in I]
     84 	b = [b[i] for i in I]
     85 
     86 	# Back substitution
     87 	for i in range(len(A)-1, -1, -1):
     88 		for j in range(i):
     89 			reducerow(A, b, i, j)
     90 
     91 	return A, b, c
     92 
     93 # Find all combinations of parameters respecting the bounds in c. The
     94 # free parameters are always the last n-rank(A) columns of the matrix A
     95 # (where n is the total number of columns).
     96 def paramcomb(nparam, c):
     97 	if nparam == 0:
     98 		return [[]]
     99 
    100 	ret = []
    101 	for i in range(c[-nparam]+1):
    102 		ret += [[i, *l] for l in paramcomb(nparam-1, c)]
    103 	return ret
    104 
    105 # Solve the integer linear system Ax = b, minimizing the sum of the
    106 # coordinates of x.
    107 def solve_system_min_sum(A, b, c):
    108 	k = len(A[0]) - len(A)
    109 	mins = inf
    110 	for c in paramcomb(k, c):
    111 		sol = sum(c)
    112 		for i in range(len(A)):
    113 			cc = (c[j]*A[i][len(A[0])-k+j] for j in range(len(c)))
    114 			s = b[i]-sum(cc)
    115 			a = s // A[i][i]
    116 			# If a is negative or not integer, we skip this solution
    117 			if a < 0 or s % A[i][i] != 0:
    118 				sol = inf
    119 				break
    120 			sol += a
    121 		mins = min(mins, sol)
    122 	return mins
    123 
    124 with fileinput.input() as lines:
    125 	print(sum(solve_system_min_sum(*reduce(*readline(l))) for l in lines))