aoc_2022_day_19_tf.py

from collections import *
import math, sys, copy, functools, json

# !!! ONE ore-collecting robot in your pack that you can use to kickstart the whole operation

# which blueprint would maximize the number of opened geodes after 24 minutes

# needs at least 16GB ram to run pt 2...

sys.setrecursionlimit(1000000000)
f = open("day19.input")
# g = f.readline()

arr = []

d = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: False)))
d = defaultdict(lambda: defaultdict(lambda: False))

bps = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: False)))


CODES = ['OR', 'C', 'OB', 'G']
CODE_TO_IDX = {
  'OR': 0,
  'C': 1,
  'OB': 2,
  'G': 3
}

def to_code(name):
  code = None
  if name == 'ore': code = 'OR'
  elif name == 'clay': code = 'C'
  elif name == 'obsidian': code = 'OB'
  elif name == 'geode': code = 'G'
  return code

MAX_BY_CODE = defaultdict(lambda: defaultdict(lambda: 0))

def insert(idx, datum):
  sp = datum.split()
  # sp.index('costs')
  code = to_code(sp[0])
  print("jaja", sp)

  costing = [z.strip() for z in " ".join(sp[sp.index('costs')+1:]).split("and")]
  # print("jaja2", sp[sp.index('costs')+1:])
  print("jaja2", costing)
  for cost in costing:
    c, name = cost.split()
    c = int(c)
    print("inserting!", (idx, code, to_code(name), c))
    bps[idx][code][to_code(name)] = c
    MAX_BY_CODE[idx][ CODE_TO_IDX[to_code(name)] ] = max( MAX_BY_CODE[idx][ CODE_TO_IDX[to_code(name)] ] ,  c )

for line in f:
  line = line.strip()
  # print(line)
  sp = line.split()
  # print(sp)
  idx = int(sp[1][:-1])
  # bps[idx] = {}

  data = [s.strip()[:-1] for s in " ".join(sp[2:]).split("Each") if s]
  for datum in data:
    insert(idx, datum)
  # print(data)

f.close()

print("bps debug")
print(json.dumps(bps[1], indent=2))
# print(bps[2]['OR'])

print(json.dumps(MAX_BY_CODE, indent=2))


res = 0

visited = set()
# earliest_time = defaultdict(lambda: math.inf)
cur_max_geode = 0
def test(idx, minutes, cur_freqs, robots):
#   if robots[0] > 4:
#     print("should never happen")
#     print(robots)
#   print("called", (idx, minutes, cur_freqs, robots))
#   for i in range(3):
#     if robots[i] > MAX_BY_CODE[i]+5:
#         return
  global res
  # if cur_freqs[-1] < res:
  #   return
  if minutes == 0:
    # print("ended", (idx, minutes, cur_freqs, robots))
    res = max(res, cur_freqs[-1]) # geodes never consumed.
    return

  # bt
  # print("hmm...", sorted(bps[idx].items(), reverse=True, key=lambda x: CODE_TO_IDX[x[0]]))
  # for code, value in bps[idx].items():
  has_big = False
  has_geode = False
  for code, value in sorted(bps[idx].items(), reverse=True, key=lambda x: CODE_TO_IDX[x[0]]):
#   for code, value in bps[idx].items():
    needed_types = len(value)
    code_idx = CODE_TO_IDX[code]

    # print("hmm?", code_idx, code)

    if code_idx < 3 and robots[code_idx] == MAX_BY_CODE[idx][code_idx]: # don't need to buy any more...
        continue
    
    n_cur_freqs = list(cur_freqs)
    n_robots = list(robots)

    for needed_code, needed_amt in value.items():
      # print( (code, needed_code, needed_amt) )
      needed_code_idx = CODE_TO_IDX[needed_code]
      if cur_freqs[needed_code_idx] >= needed_amt:
        needed_types -= 1
        # n_cur_freqs = cur_freqs[:]
        n_cur_freqs[needed_code_idx] -= needed_amt
        # n_robots = robots[:]
    if (needed_types == 0):
      n_robots[code_idx] += 1
      if has_geode: continue
    #   if has_big and code_idx < 2: continue
      # bought = True
      # print("appending!", ( total_opened[:], n_cur_freqs, n_robots) )
      for i in range(len(robots)):
        n_cur_freqs[i] += robots[i]
      
      hsh = (minutes-1, *(n_cur_freqs), *(n_robots) )
    #   if earliest_time[hsh] > minutes:
        # earliest_time[hsh] = minutes
      if hsh not in visited:
        visited.add(hsh)
        test(idx, minutes-1, tuple(n_cur_freqs), tuple(n_robots))
      if code_idx == 3:
        has_big = True
        has_geode = True
      if code_idx == 2:
        has_big = True
        # break
      # else:
      #   print("wtf?")
  if has_geode: return  
  n_cur_freqs = list(cur_freqs)
  for i in range(len(robots)):
    n_cur_freqs[i] += robots[i]
  
  hsh = (minutes-1, *(n_cur_freqs), *(robots) )
#   if earliest_time[hsh] > minutes:
#     earliest_time[hsh] = minutes
  if hsh not in visited:
    visited.add(hsh)
    test(idx, minutes-1, tuple(n_cur_freqs), robots)
  # else:
  #   print("wtf 2?")
  # test(idx, minutes-1, tuple(n_cur_freqs), robots)


zres = 0
for idx in list(bps.keys()):
  res = 0
  visited.clear()
  test(idx, 32, (0,0,0,0), (1,0,0,0) )
  print("tested", idx, res)
  zres += idx * res
  # break
print("final", zres)