advent2020/twentytwo.py

102 lines
2.7 KiB
Python
Executable File

#!/usr/bin/env python3
from functools import reduce
import datetime
import re, sys
import math
import numpy as np
from numba import jit
def tpl(l):
return int(l)
def dprint(*args):
sep = False
for arg in args:
if sep:
print('--------')
sep = True
for a, b in arg.items():
print('%10s: %s' % (a, ','.join(str(x) for x in b)))
print()
@jit
def rnd(carda, cardb, idx):
if idx == 0:
return [carda, cardb] if carda > cardb else []
if idx == 1:
return [cardb, carda] if carda < cardb else []
@jit
def rnd2(carda, cardb, idx, winner):
if idx == 0:
return [carda, cardb] if idx == winner else []
if idx == 1:
return [cardb, carda] if idx == winner else []
@jit
def hashish(fp, sp):
return (reduce(lambda a, b: a|(1<<b), fp, 0) << 50) + \
reduce(lambda a, b: a|(1<<b), sp, 0)
@jit
def finalcount(pile):
return sum([i * x for i, x in enumerate([0, ] + pile[::-1])])
@jit
def game(pilea, pileb, rec=1):
#print('Game %d' % rec)
track = {}
rd = 1
tot = len(pilea) + len(pileb)
while len(pilea) not in (0, tot):
if hashish(pilea,pileb) in track:
return 0, -1
#print('Game %d round %d' % (rec, rd))
track[hashish(pilea,pileb)] = True
#dprint({'1': pilea, '2':pileb})
a, b = pilea[0], pileb[0]
if len(pilea[1:]) >= a and len(pileb[1:]) >= b:
winner, _ = game(pilea[1:a+1], pileb[1:b+1], rec+1)
pilea = pilea[1:] + rnd2(a, b, 0, winner)
pileb = pileb[1:] + rnd2(a, b, 1, winner)
else:
pilea = pilea[1:] + rnd(a, b, 0)
pileb = pileb[1:] + rnd(a, b, 1)
rd += 1
if len(pileb):
if rec > 1:
return 1, -1
else:
return 1, finalcount(pileb)
else:
if rec > 1:
return 0, -1
else:
return 0, finalcount(pilea)
fn = sys.argv[1] if len(sys.argv) > 1 else 'input%s' % 22
with open(fn) as f:
#pgphs = f.read().split('\n\n')
header, footer = f.read().split('\n\n')
p1 = [tpl(l) for l in header.strip().split('\n')[1:]]
p2 = [tpl(l) for l in footer.strip().split('\n')[1:]]
tot = len(p1) + len(p2)
while len(p1) not in (0, tot):
#dprint({'1': p1, '2':p2})
a, b = p1[0], p2[0]
p1 = p1[1:] + rnd(a, b, 0)
p2 = p2[1:] + rnd(a, b, 1)
final = sum(i * x for i, x in enumerate([0, *p1[::-1]]))
#print(final)
final += sum(i * x for i, x in enumerate([0, *p2[::-1]]))
print(final)
p1 = [tpl(l) for l in header.strip().split('\n')[1:]]
p2 = [tpl(l) for l in footer.strip().split('\n')[1:]]
print(game(p1, p2)[1])