aoc18-4.py

with open('input') as f:
  ss = f.read().splitlines()
  
sorted_input=sorted(ss,key=lambda d:d[1:17])

bucket=[]
res=[]

#partition by begin shift message
for entry in sorted_input:
  if '#' in entry and not bucket:
    bucket.append(entry)
  elif bucket and '#' not in entry:
    bucket.append(entry)
  elif bucket and '#' in entry:
    res.append(bucket)
    bucket = [entry]
if bucket:
  res.append(bucket)

#process each shift
import re
id_pattern = r'#\d+'
minute_pattern_sleep = r'\d+(?=] falls asleep)' 
minute_pattern_wake = r'\d+(?=] wakes up)'
times = {}

def process_shift(shift):
  #exploit structure that id is first then sleep wake alternates
  id = re.findall(id_pattern, shift[0])[0]
  sleep = re.findall(minute_pattern_sleep, ','.join(shift[1::2]))
  wake = re.findall(minute_pattern_wake, ','.join(shift[2::2]))
  
  assert(len(sleep) == len(wake))

  try:
    for (start,end) in zip(sleep,wake):
      for i in range(int(start),int(end)):
        times[id][i] += 1
  except KeyError:
    #initialize it if it's not there
    times[id] = [0 for i in range(61)]

#apply the processing
list(map(lambda x: process_shift(x), res))

#id : minutes asleep, minute most asleep
total_sleeps = {k:(sum(v),v.index(max(v))) for (k,v) in times.items()}
ans_1= max(total_sleeps.items(),key=lambda x:x[1])

#part 2
most_asleep_minutes = {k:max(enumerate(v), key = lambda x:x[1]) for (k,v) in times.items()}
ans_2=max(most_asleep_minutes.items(),key=lambda x:x[-1][-1])