aoc/2023/13/solution.py

#!/usr/bin/python3
import sys

def pattern_str(pattern):
    return '\n'.join([''.join(line) for line in pattern])

def rotate(pattern):
    return [list(line) for line in zip(*reversed(pattern))]

def is_reflected(up, down, k):
    for l in range(min(len(up), len(down))):
        i, j = k - l - 1, l
        if up[i] != down[j]:
            return False
    return True

def get_reflection(pattern):
    total = len(pattern)
    for k in range(1, len(pattern)):
        print(f'k={k}')
        if is_reflected(pattern[:k], pattern[k:], k):
            return k
    return -1

def is_reflected_with_correction(up, down, k):
    correction = True
    for l in range(min(len(up), len(down))):
        i, j = k - l - 1, l
        if up[i] != down[j] and correction:
            different = [(u, d) for u, d in zip(up[i], down[j]) if u != d]
            if len(different) == 1:
                #print(different, len(different))
                correction = False
                continue
            else:
                return False
        elif up[i] != down[j]:
            return False
    return not correction

def get_reflection_with_correction(pattern):
    total = len(pattern)
    max_k = -1
    for k in range(1, len(pattern)):
        #print(f'k={k}')
        if is_reflected_with_correction(pattern[:k], pattern[k:], k):
            max_k = max(max_k, k)
    return max_k

def part12(correction=False):
    print('part 1')
    patterns = [[]]
    for line in map(lambda line: line.strip(), sys.stdin):
        if len(line) == 0:
            patterns.append([])
        else:
            patterns[-1].append(list(line))
    total = 0
    for i, pattern in enumerate(patterns):
        print(f'pattern {i}:\n{pattern_str(pattern)}\nrotated:\n{pattern_str(rotate(pattern))}')
        reflection = get_reflection_with_correction(pattern) if correction else get_reflection(pattern)
        if reflection > 0:
            print(f'found horizontal reflection at {reflection}')
            total += 100 * reflection
            continue
        rotated = rotate(pattern)
        reflection = get_reflection_with_correction(rotated) if correction else get_reflection(rotated)
        if reflection > 0:
            print(f'found vertical reflection at {reflection}')
            total += reflection
    print(total)

if sys.argv[1] in '1':
    part12(correction=False)
else:
    part12(correction=True)