Python Picross solver

picross-solver.py

import functools
import itertools
import operator

@functools.lru_cache(maxsize=None)
def generate_theories(constraint, length):
    if 0 == len(constraint):
        return [(False,) * length]
    excess = length - len(constraint) + 1 - sum(constraint)
    theories = []
    for space in range(excess+1):
        prefix = (False,) * space + (True,) * constraint[0]
        remaining = length-constraint[0]-space
        if 1 < len(constraint):
            prefix += (False,)
            remaining -= 1
        for t in generate_theories(constraint[1:], remaining):
            theories.append(prefix + t)
    return theories

def find_theory_overlaps(theories):
    assert(len(theories))
    overlap = []
    for guesses in zip(*theories):
        if all(guesses):
            overlap.append(True)
        elif not any(guesses):
            overlap.append(False)
        else:
            overlap.append(None)
    return overlap

def filter_theories(overlap, theories):
    bad_theories = set()
    for idx, v in enumerate(overlap):
        if v is None:
            continue
        guesses = map(operator.itemgetter(idx), theories)
        if v:
            bad_guesses = itertools.filterfalse(operator.itemgetter(1),
                                                enumerate(guesses))
        else:
            bad_guesses = filter(operator.itemgetter(1),
                                 enumerate(guesses))
        bad_theories.update(map(operator.itemgetter(0), bad_guesses))
    return [t
            for idx,t in enumerate(theories)
            if idx not in bad_theories]

def generate_grid(scale):
    return [[None for _ in range(scale[1])]
            for _ in range(scale[0])]

def solve(row_constraints, column_constraints, scale, progress=None):
    row_theories = [generate_theories(tuple(r), scale[0])
                    for r in row_constraints]
    column_theories = [list(generate_theories(tuple(c), scale[1]))
                       for c in column_constraints]

    generate_theories.cache_clear()

    grid = generate_grid(scale)

    rows_to_check = range(scale[0])
    cols_to_check = range(scale[1])

    iteration = 0
    while True:
        iteration += 1

        if progress:
            progress.start_iteration(iteration)
            progress.unknown_cells(sum(1
                                       for row in grid
                                       for v in row
                                       if v is None))

        cols_to_update = set()
        for row_idx in sorted(rows_to_check):
            overlap = find_theory_overlaps(row_theories[row_idx])
            for col_idx, v in enumerate(overlap):
                if v is None:
                    continue
                grid_val = grid[row_idx][col_idx]
                if grid_val == v:
                    continue
                assert(grid_val is None)
                grid[row_idx][col_idx] = v
                cols_to_update.add(col_idx)

        rows_to_update = set()
        for col_idx in sorted(cols_to_check):
            overlap = find_theory_overlaps(column_theories[col_idx])
            for row_idx, v in enumerate(overlap):
                if v is None:
                    continue
                grid_val = grid[row_idx][col_idx]
                if grid_val == v:
                    continue
                assert(grid_val is None)
                grid[row_idx][col_idx] = v
                rows_to_update.add(row_idx)

        if progress:
            progress.new_grid(grid)

        rows_to_check = []
        for idx in sorted(rows_to_update):
            rt = row_theories[idx]
            overlap = grid[idx][:]
            new_theories = filter_theories(overlap, rt)
            assert(len(new_theories))
            row_theories[idx] = new_theories
            if len(new_theories) < len(rt):
                rows_to_check.append(idx)

        cols_to_check = []
        for idx in sorted(cols_to_update):
            ct = column_theories[idx]
            overlap = [row[idx]
                       for row in grid]
            new_theories = filter_theories(overlap, ct)
            assert(len(new_theories))
            column_theories[idx] = new_theories
            if len(new_theories) < len(ct):
                cols_to_check.append(idx)

        if not any(v == None
                   for row in grid
                   for v in row):
            if progress:
                progress.solve_complete()
            return grid

        assert(rows_to_check or cols_to_check)

        if progress:
            progress.end_iteration()