Python
Since, an optimized phase 2 is enough to solve p2 and p3, I did not optimize any further.
import math
# Brute-force phase one
# Simulate all phase one rounds until no more moves can be made or max rounds reached
def phase_one(ducks: list[int], max_rounds: int) -> int:
n = len(ducks)
round = 1
while round <= max_rounds:
moved = False
for i in range(n - 1):
if ducks[i] > ducks[i + 1]:
ducks[i] -= 1
ducks[i + 1] += 1
moved = True
if not moved:
break
round += 1
return round - 1
# Brute-force phase two
# Simulate all phase two rounds until no more moves can be made or max rounds reached
def phase_two(ducks: list[int], max_rounds: int) -> int:
n = len(ducks)
round = 1
while round <= max_rounds:
moved = False
for i in range(n - 1):
if ducks[i] < ducks[i + 1]:
ducks[i] += 1
ducks[i + 1] -= 1
moved = True
if not moved:
break
round += 1
return round - 1
# Optimized phase two for limitless rounds
# Every round will move one duck from every duck above average to below average
# So the resulting number of rounds is simply the total number of ducks above average
# Inversely, you can also total the number of ducks below average
def phase_two_limitless(ducks: list[int]) -> int:
avg = 0
for i, d in enumerate(ducks):
avg += (d - avg) / (i + 1)
avg = round(avg)
rounds = sum(x - avg for x in ducks if x > avg)
return rounds
# Balance ducks through both phases, respecting max rounds
def balance_ducks(ducks: list[int], max_rounds: int) -> int:
rounds1 = phase_one(ducks, max_rounds)
# if the number of rounds is unbounded, use the optimized phase two
if max_rounds < math.inf:
rounds2 = phase_two(ducks, max_rounds - rounds1)
else:
rounds2 = phase_two_limitless(ducks)
return rounds1 + rounds2
# Calculate checksum of the ducks
def calculate_checksum(ducks: list[int]) -> int:
n = len(ducks)
checksum = sum((i + 1) * ducks[i] for i in range(n))
return checksum
# Brute-force both phases
def part1(data: str):
ducks = [int(d) for d in data.splitlines()]
balance_ducks(ducks, 10)
return calculate_checksum(ducks)
# <asserts snipped>
# Brute-force phase one, use optimized phase two
def part2(data: str):
ducks = [int(d) for d in data.splitlines()]
rounds = balance_ducks(ducks, math.inf) # type: ignore
return rounds
# <asserts snipped>
# You can use the same function for p2 and p3
part3 = part2