50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
from tqdm import tqdm
|
|
from argparse import ArgumentParser
|
|
import sys
|
|
|
|
sys.setrecursionlimit(10**7)
|
|
|
|
def transform(num_str: str) -> list:
|
|
# Apply one iteration of the rules
|
|
if num_str == '0':
|
|
return ["1"]
|
|
elif len(num_str) % 2 == 0:
|
|
half = len(num_str) // 2
|
|
return [str(int(num_str[:half])), str(int(num_str[half:]))]
|
|
else:
|
|
return [str(int(num_str) * 2024)]
|
|
|
|
def expand(num: str, iteration: int, cache: dict) -> int:
|
|
# Returns count of stones after 'iteration' blinks
|
|
if iteration == 0:
|
|
return 1
|
|
key = (num, iteration)
|
|
if key in cache:
|
|
return cache[key]
|
|
nxt = transform(num)
|
|
total = 0
|
|
for n in nxt:
|
|
total += expand(n, iteration - 1, cache)
|
|
cache[key] = total
|
|
return total
|
|
|
|
def main(input_str: str, iterations: int) -> int:
|
|
stones = input_str.split(" ")
|
|
cache = {}
|
|
count = 0
|
|
with tqdm(total=len(stones), desc="Expanding stones") as pbar:
|
|
for s in stones:
|
|
count += expand(s, iterations, cache)
|
|
pbar.update(1)
|
|
return count
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser()
|
|
parser.add_argument("-p", "--part")
|
|
parser.add_argument("-f", "--file")
|
|
args = parser.parse_args()
|
|
|
|
inp = open(args.file, 'r').read().strip()
|
|
iters = 25 if args.part == "1" else (75 if args.part == "2" else None)
|
|
|
|
print(main(inp, iters))
|