aoc/2023/11/solution.py

#!/usr/bin/python3
import sys

def expand_system(system, m, factor):
    empty_rows = set(range(m)) - {y for y, _ in system}
    empty_cols = set(range(m)) - {x for _, x in system}
    expanded = set()
    for y, x in system:
        expanded.add((y + sum(row < y for row in empty_rows) * factor, (x + sum(col < x for col in empty_cols) * factor)))
    return expanded

def manhattan(a, b):
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def part12(factor=1):
    print('part 1')
    system = set([])
    max_y = 0
    max_x = 0
    for y, line in enumerate(sys.stdin):
        for x, c in enumerate(line.strip()):
            if c != '.':
                system.add((y, x))
    max_y = y
    max_x = x

    print(f'system ({max_y}x{max_x}) {system}')

    assert max_y == max_x
    m = max_y

    distances = 0
    system = expand_system(system, m, factor)
    for galaxy1 in system:
        for galaxy2 in system:
            if galaxy1 == galaxy2:
                continue
            distances += manhattan(galaxy1, galaxy2)
    print(distances // 2)

if sys.argv[1] in '1':
    part12(factor=1)
else:
    part12(factor=(1000000 - 1))