Part 1: Given a list of seating preferences of the form Alice would gain 54 happiness units by sitting next to Bob. find the seating arrangement which maximizes total happiness.
happiness = collections.defaultdict(lambda : collections.defaultdict(lambda : 0))
for line in sys.stdin:
self, gain_lose, amount, other = re.match(
r'(\w+) would (gain|lose) (\d+) happiness units by sitting next to (\w+).',
line
).groups()
amount = int(amount)
if gain_lose == 'lose':
amount *= -1
happiness[self][other] = amount
best_score, best_ordering = max(
(
sum(
happiness[a][b] + happiness[b][a]
for a, b in zip(ordering, ordering[1:] + (ordering[0],))
),
ordering
)
for ordering in itertools.permutations(happiness.keys())
)
print(best_ordering)
print(best_score)
I think I may have over functional’ed this one. 😄 Basically, we parse everything first into a nested dictionary of happiness scores. Since we always have to add both directions in, it doesn’t matter which is which in the indicies so long as it’s consistent.
After that, we’ll use the same trick as in Advent of Code: Day 9: calculate pairs of scores and orderings and find the largest one. To specifically calculate the total score, we add the happiness in both directions for every pair of neighbors once again using zip. The only difference is that we add (ordering[0],) back onto the end so that we don’t lose the first/last pairing of the list.
Part 2: Add one more person me, who has a mutual happiness score of ±0 with every other person.
There’s not actually terribly much different here, since the happiness structure already contains a default of 0 for any missing user. Because of that, we can just change this line:
for ordering in itertools.permutations(happiness.keys())
to this:
for ordering in itertools.permutations(list(happiness.keys()) + ['me'])
It takes a bit longer since the number of permutations grows rather quickly, but still not more than a few seconds total runtime.
Shiny.