Logic Traces Solver

logic-traces.py

#!/usr/bin/env python
from __future__ import print_function
"""
Usage example:
input via stdin like:
4 . . . . . . 3 .
. . . . . . . . 7
. 5 . . . . 3 . .
. . . 7 . . . . .
. 5 . . . . . . 7
2 . 3 . . 2 . . .
. . . . . . 7 . .
. . 3 . . . . 2 .
. 6 . . . . . . .

output via stdout like:
4 > > > < < < 3 ^
v ^ < < < < < < 7
< 5 > > > < 3 > >
< < < 7 > > > > ^
< 5 > > > > < < 7
2 < 3 > > 2 > > v
v < < < < < 7 ^ v
v ^ 3 > > > v 2 v
< 6 > > > > v v v
"""
import re
from collections import defaultdict

import pulp

class NameGenerator(object):
    def __init__(self, prefix='x'):
        self.prefix = prefix
        self.current = 0

    def generate(self):
        self.current += 1
        name = self.prefix + str(self.current)
        return name

generate_name = NameGenerator().generate

def Variable(lbound, ubound):
    return pulp.LpVariable(generate_name(), lbound, ubound, pulp.LpInteger)

def BinaryVariable():
    return Variable(0, 1)


def value(variable):
    v = pulp.value(variable)
    if v is None:
        if variable.getLb() == variable.getUb():
            return variable.getLb()
        return v
    if v == int(v):
        return int(v)
    return v


LEFT = (-1, 0)
UP = (0, -1)
RIGHT = (1, 0)
DOWN = (0, 1)

def preprocess(board):
    width = len(board[0])
    height = len(board)
    traces = {}
    for y, line in enumerate(board):
        for x, cell in enumerate(line):
            if cell is not None:
                traces[x, y] = cell

    branch_limits = defaultdict(dict)

    for pos in traces:
        total_length = traces[pos]
        for direction in [LEFT, UP, RIGHT, DOWN]:
            limit = 0
            current_pos = pos
            while limit < total_length:
                current_pos_x, current_pos_y = current_pos
                next_pos_x = current_pos_x + direction[0]
                next_pos_y = current_pos_y + direction[1]
                next_pos = (next_pos_x, next_pos_y)
                next_pos_x_in_board = 0 <= next_pos_x < width
                next_pos_y_in_board = 0 <= next_pos_y < height
                next_pos_in_board = next_pos_x_in_board and next_pos_y_in_board
                if not next_pos_in_board or next_pos in traces:
                    break
                current_pos = next_pos
                limit += 1
            if limit > 0:
                branch_limits[pos][direction] = limit

    constrains = []

    branch_lengths = defaultdict(dict)
    branch_existences = defaultdict(dict)
    branch_direction_choices = defaultdict(dict)

    for pos in branch_limits:
        total_length = traces[pos]
        for direction in branch_limits[pos]:
            limit = branch_limits[pos][direction]
            branch_length = Variable(0, limit)
            branch_lengths[pos][direction] = branch_length
            branch_exists = BinaryVariable()
            branch_existences[pos][direction] = branch_exists

            constrains.append(branch_exists * limit >= branch_length)
            constrains.append(branch_exists <= branch_length)

            current_pos = pos
            for i in range(1, limit + 1):
                current_pos_x, current_pos_y = current_pos
                next_pos_x = current_pos_x + direction[0]
                next_pos_y = current_pos_y + direction[1]
                next_pos = (next_pos_x, next_pos_y)
                current_pos = next_pos
                direction_choice = BinaryVariable()
                branch_direction_choices[current_pos][direction] = direction_choice
                # direction_choice == 1 iff i <= branch_length
                # direction_choice * i <= branch_length
                # (1 - direction_choice) * (limit + 1 - i) <= limit - branch_length
                constrains.append(direction_choice * i <= branch_length)
                constrains.append((1 - direction_choice) * (limit + 1 - i) <= limit - branch_length)

        constrains.append(pulp.lpSum(branch_lengths[pos].values()) == total_length)

    for pos in branch_direction_choices:
        constrains.append(pulp.lpSum(branch_direction_choices[pos].values()) == 1)

    objective = 0
    objective_limit = 0

    for pos in branch_existences:
        for direction in branch_existences[pos]:
            objective += branch_existences[pos][direction]
            objective_limit += 1

    objective_var = Variable(0, objective_limit)
    constrains.append(objective_var == objective)

    prob = pulp.LpProblem('Logic Traces Problem', pulp.LpMinimize)
    prob += objective_var
    for constrain in constrains:
        prob += constrain

    return prob, (width, height, traces, branch_direction_choices)


def postprocess(namespace):
    direction_mapping = {
        LEFT: '<',
        UP: '^',
        RIGHT: '>',
        DOWN: 'v',
    }
    width, height, traces, branch_direction_choices = namespace
    cell_width = len(str(max(traces.values())))
    for y in range(height):
        line = []
        for x in range(width):
            if (x, y) in traces:
                cell = str(traces[x, y])
            else:
                direction_vars = branch_direction_choices[x, y]
                direction = max(direction_vars, key=lambda d:value(direction_vars[d]))
                cell = direction_mapping[direction]
            line.append(cell.rjust(cell_width))
        print(' '.join(line))


def read_board(input):
    board = []
    while True:
        raw_line = input().strip()
        if not raw_line:
            break
        line = []
        for match in re.finditer(r'(\.)|(\d+)', raw_line):
            if match.group(1):
                line.append(None)
            else:
                line.append(int(match.group(2)))
        board.append(line)
    return board


try:
    _input = raw_input
except NameError:
    _input = input
board = read_board(_input)
prob, namespace = preprocess(board)
if prob.solve() == pulp.LpStatusOptimal:
    postprocess(namespace)