kenken/main.py

import itertools
from copy import deepcopy
from itertools import combinations, combinations_with_replacement

#             up       right   down    left
DIRECTIONS = [(0, -1), (1, 0), (0, 1), (-1, 0)]
OP_PLUS = "+"
OP_MINUS = "-"
OP_MULTIPLY = "x"
OP_DIVIDE = "/"
OP_NONE = "#"

BOX_H_LABELS = 'ABCDEFGHI'
BOX_V_LABELS = '123456789'


def lambda_mul(x):
    rv = 1
    for v in x:
        rv *= v
    return rv


OP_LAMBDAS = {
    OP_PLUS: sum,
    OP_MINUS: lambda x: abs(x[0] - x[1]),
    OP_MULTIPLY: lambda_mul,
    OP_DIVIDE: lambda x: x[0] / x[1] if x[0] % x[1] == 0 else x[1] / x[0],
    OP_NONE: lambda x: x[0],
}


def translate_box_to_rc(box, grid_size):
    return BOX_V_LABELS.index(box[1]), BOX_H_LABELS[grid_size - 1::-1].index(box[0])


def translate_rc_to_box(row_id, col_id, grid_size):
    return BOX_H_LABELS[col_id] + BOX_V_LABELS[grid_size - row_id - 1]


def has_line_integrity(line):
    existing_values = set()
    for value in line:
        if value in existing_values:
            return False
        if value != 0:
            existing_values.add(value)
    return True


def fill_grid(grid, solution):
    box_values = []
    for box_value in solution:
        box_values.append(box_value)
    for box, value in box_values:
        br, bc = translate_box_to_rc(box, len(grid))
        grid[br][bc] = value
    return grid


def check_grid(grid):
    for row in grid:
        if not has_line_integrity(row):
            return False
    for col_id in range(len(grid)):
        col = [grid[r][col_id] for r in range(len(grid))]
        if not has_line_integrity(col):
            return False
    return True


def init_grid(grid_size):
    return [[0] * grid_size for _ in range(grid_size)]


class Block(object):
    def __init__(self, boxes, operation, result):
        self.boxes = boxes
        self.operation = operation
        self.result = result

        self.solutions = set()

        self.verify()

    def verify(self):
        """
        Verifies that the box location are valid in the range A-F and 1-6
        Operations are in +-*/ or None.
        Where -/ have exactly 2, None has exactly 1.
        """
        assert self.operation in OP_LAMBDAS
        if self.operation in [OP_NONE]:
            assert len(self.boxes) == 1
        elif self.operation in [OP_MINUS, OP_DIVIDE]:
            assert len(self.boxes) == 2
        else:
            assert len(self.boxes) > 1
        for box in self.boxes:
            assert len(box) == 2
            assert box[0] in BOX_H_LABELS
            assert box[1] in BOX_V_LABELS

    def all_boxes_in_one_row_or_column(self):
        """Returns true if all of the boxes are in the same row OR the same column"""
        is_same_column, is_same_row = True, True
        for box in self.boxes[1:]:
            if self.boxes[0][0] != box[0]:
                is_same_column = False
            if self.boxes[0][1] != box[1]:
                is_same_row = False
        return is_same_row or is_same_column

    def generate_combinations(self, k, n):
        """Generates k|n combinations for 2 boxes or if they are in the same row OR column.
        Generates k|n combinations with replacements for other cases."""
        assert 1 <= k < n
        # Exploit the structure of the problem
        # For block of size 2, we don't need comb with replacement as they will
        # neccesarily be in different columns/rows
        if k == 2:
            return list(combinations(range(1, n + 1), k))
        # for 3 and more we don't need if all of the boxes are on the
        # same row or same column
        elif k > 2:
            if self.all_boxes_in_one_row_or_column():
                return list(combinations(range(1, n + 1), k))
        return list(combinations_with_replacement(range(1, n + 1), k))

    def generate_hypotheses(self, grid_size):
        """Generates hypothesis for this block based on the operation and result."""
        rv = []
        op = OP_LAMBDAS[self.operation]
        hypotheses = self.generate_combinations(k=len(self.boxes), n=grid_size)
        for hypothesis in hypotheses:
            if op(hypothesis) == self.result:
                rv.append(hypothesis)
        return rv

    def generate_solutions(self, grid_size):
        """Generates possible solutions for the block, including the limiting row/column requirement."""
        self.generate_hypotheses(grid_size)
        for hyp in self.generate_hypotheses(grid_size):
            for perm in itertools.permutations(hyp):
                sol = tuple(zip(self.boxes, perm))
                if check_grid(fill_grid(init_grid(grid_size), sol)):
                    self.solutions.add(sol)

    def __repr__(self):
        return 'Block {}'.format(self.boxes)


class Game(object):
    def __init__(self, grid_size: int, blocks):
        assert 2 <= grid_size <= 9
        self.blocks = blocks
        self.grid_size = grid_size
        self.verify()

        for block in self.blocks:
            block.generate_solutions(self.grid_size)

    def verify(self):
        """Check that each block is found exactly once in the grid."""
        found = set()
        for block in self.blocks:
            for box in block.boxes:
                assert box not in found
                found.add(box)
        assert len(found) == self.grid_size ** 2

    def solve(self, grid, current_block_id=0):
        """
        Recursive solution - Start by filling up the first block.
        For each solution that passes the row/column requirement, check the recursive solution
        with the next block. Break when it's the last block and the row/col requirement is filled.
        """
        block = self.blocks[current_block_id]
        for solution in block.solutions:
            prev_grid = deepcopy(grid)
            grid = fill_grid(grid, solution)
            if not check_grid(grid):
                grid = prev_grid
            else:
                if current_block_id == len(self.blocks) - 1:
                    return grid
                sol = self.solve(deepcopy(grid), current_block_id + 1)
                if sol:
                    return sol


def walk_grid_to_boxes(grid, row_id, col_id, initial_cell, boxes):
    """Recursively walk through the boxes in a grid going up,left,down,right to construct
    a contiguous block of same operations."""
    grid_size = len(grid)
    for v_offset, h_offset in DIRECTIONS:
        new_row_id, new_col_id = row_id + h_offset, col_id + v_offset
        # avoid wrap-around
        if not (0 <= new_row_id < grid_size and 0 <= new_col_id < grid_size):
            continue
        box = translate_rc_to_box(new_row_id, new_col_id, grid_size)
        if box in boxes:
            continue
        cell = grid[new_row_id][new_col_id].strip()
        if cell == initial_cell:
            boxes.append(box)
            walk_grid_to_boxes(grid, new_row_id, new_col_id, initial_cell, boxes)
    return boxes


def construct_blocks(grid):
    """Reads a grid of operations and converts them to contiguous blocks."""
    grid_size = len(grid)
    blocks, used_boxes = [], []
    for row_id, row in enumerate(grid):
        assert len(row) == grid_size
        for col_id, cell in enumerate(row):
            box = translate_rc_to_box(row_id, col_id, grid_size)
            if box in used_boxes:
                continue
            cell = cell.strip()
            operation = cell[-1]
            assert operation in OP_LAMBDAS.keys()
            value = int(cell[:-1])
            boxes = walk_grid_to_boxes(grid, row_id, col_id, cell, [box])
            used_boxes += boxes
            blocks.append(Block(boxes=boxes, operation=operation, result=value))
    return grid_size, blocks


def print_grid(grid):
    for line in grid:
        print(line)


def solve_game(grid: list):
    grid_size, blocks = construct_blocks(grid)
    game = Game(grid_size=grid_size, blocks=blocks)
    grid = game.solve(grid=init_grid(game.grid_size))
    print_grid(grid)


def main():
    """Construct a game using the +-x/ symbols for operations and the number before that.
    Whitespace around is meaningless and can be used to arrange the matrix visually."""
    solve_game([
        ['120x', '120x', '120x', '120x', '17+', '5+ '],
        ['1-  ', '1-  ', '2/  ', '3+  ', '17+', '5+ '],
        ['5+  ', '30x ', '2/  ', '3+  ', '17+', '17+'],
        ['5+  ', '30x ', '30x ', '15x ', '16+', '16+'],
        ['6#  ', '10x ', '10x ', '15x ', '15x', '16+'],
        ['3/  ', '3/  ', '10x ', '1-  ', '1- ', '16+'],
    ])

    # Alternatively construct a Game object more explicitly:
    # game = Game(grid_size=6,
    #             blocks=[
    #                 Block(boxes=['A6', 'B6', 'C6', 'D6'], operation=OP_MULTIPLY, result=120),
    #                 Block(boxes=['E6', 'E5', 'E4', 'F4'], operation=OP_PLUS, result=17),
    #                 Block(boxes=['F6', 'F5'], operation=OP_PLUS, result=5),
    #                 Block(boxes=['A5', 'B5'], operation=OP_MINUS, result=1),
    #                 Block(boxes=['C5', 'C4'], operation=OP_DIVIDE, result=2),
    #                 Block(boxes=['D5', 'D4'], operation=OP_PLUS, result=3),
    #                 Block(boxes=['A4', 'A3'], operation=OP_PLUS, result=5),
    #                 Block(boxes=['B4', 'B3', 'C3'], operation=OP_MULTIPLY, result=30),
    #                 Block(boxes=['D3', 'D2', 'E2'], operation=OP_MULTIPLY, result=15),
    #                 Block(boxes=['B2', 'C2', 'C1'], operation=OP_MULTIPLY, result=10),
    #                 Block(boxes=['A2'], operation=OP_NONE, result=6),
    #                 Block(boxes=['A1', 'B1'], operation=OP_DIVIDE, result=3),
    #                 Block(boxes=['D1', 'E1'], operation=OP_MINUS, result=1),
    #                 Block(boxes=['F1', 'F2', 'F3', 'E3'], operation=OP_PLUS, result=16),
    #             ])


if __name__ == '__main__':
    main()