day9.py

import re
from functools import *


def read_file(filename):
    with open(filename) as f:
        return f.readline()


def decompressed_string(str, version=1):
    repeat_re = re.compile(r'\((\d+)x(\d+)\)')
    current_idx = 0

    while True:
        match = repeat_re.search(str[current_idx:])
        if match:
            if match.span()[0] > 0:
                yield str[current_idx:current_idx+match.span()[0]]
            current_idx += match.span()[1]
            count = int(match[1])
            for n in range(0, int(match[2])):
                if version == 1:
                    yield str[current_idx:current_idx+count]
                else:
                    for s in decompressed_string(str[current_idx:current_idx+count], 2):
                        yield s
            current_idx += int(match[1])
        else:
            yield str[current_idx:]
            break


def decompressed_len(str, current_idx = 0, end_idx = None):
    repeat_re = re.compile(r'\((\d+)x(\d+)\)')
    end_idx = end_idx or len(str)

    while True:
        match = repeat_re.search(str[current_idx:end_idx])
        if match:
            if match.span()[0] > 0:
                yield match.span()[0]
            current_idx += match.span()[1]
            count = int(match[1])
            sublen = reduce(lambda a,b:a+b, decompressed_len(str, current_idx, current_idx+count))
            yield int(match[2])*sublen
            current_idx += int(match[1])
        else:
            yield end_idx-current_idx
            break

print('Part 1: %d' % reduce(lambda a, b: a+len(b), decompressed_string(read_file('day9.txt')), 0))
print('Part 2: %d' % reduce(lambda a, b: a+b, decompressed_len(read_file('day9.txt'))))