-
Notifications
You must be signed in to change notification settings - Fork 0
/
kenken_solver.py
124 lines (97 loc) · 3.41 KB
/
kenken_solver.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
# A solver for KenKen puzzles
from collections import defaultdict
from itertools import chain, product
from constraint import AllDifferentConstraint, FunctionConstraint, Problem
def AddsTo(total):
return lambda *args: sum(args) == total
def MultipliesTo(total):
return lambda *args: reduce(lambda x, y: x * y, args) == total
def DividesTo(total):
def fun(*args):
args = list(args[:])
args.sort(reverse=True)
result = args.pop(0)
for arg in args:
result /= arg
return result == total
return fun
def SubtractsTo(total):
return lambda *args: (
reduce(lambda x, y: max(x, y) - min(x, y), args) == total)
def solve_board(board, cages):
problem = Problem()
height = len(board)
width = len(board[0])
assert width == height, 'Grid must be a square'
cage_name_to_locations = defaultdict(list)
for x in range(height):
row_variables = [(x, ry) for ry in range(width)]
column_variables = [(cx, x) for cx in range(height)]
problem.addConstraint(AllDifferentConstraint(), row_variables)
problem.addConstraint(AllDifferentConstraint(), column_variables)
for y in range(width):
if isinstance(board[x][y], basestring):
# we are dealing with a function
cage_name_to_locations[board[x][y]].append((x, y))
else:
# we are dealing with a pre-assigned number
problem.addVariable((x, y), [board[x][y]])
for cage_name, cage_locations in cage_name_to_locations.iteritems():
cage_function = cages[cage_name]
all_values = product(range(1, width + 1),
repeat=len(cage_locations))
possible_values = set(chain(*[values for values in all_values
if cage_function(*values)]))
for location in cage_locations:
problem.addVariable(location, list(possible_values))
problem.addConstraint(FunctionConstraint(cage_function),
cage_locations)
solution = problem.getSolution()
answer = [row[:] for row in board]
for x in range(height):
for y in range(width):
answer[x][y] = solution[(x, y)]
return answer
board = [['a', 'a', 2],
['b', 'c', 'c'],
['b', 'd', 'd']]
cages = {
'a': SubtractsTo(2),
'b': DividesTo(2),
'c': DividesTo(3),
'd': SubtractsTo(1)
}
answer = solve_board(board, cages)
assert answer == [[3, 1, 2],
[2, 3, 1],
[1, 2, 3]]
board = [['a', 'b', 'b', 'c', 'd', 'd'],
['a', 'e', 'e', 'c', 'f', 'd'],
['g', 'g', 'h', 'h', 'f', 'd'],
['g', 'g', 'i', 'j', 'k', 'k'],
['l', 'l', 'i', 'j', 'j', 'm'],
['n', 'n', 'n', 'o', 'o', 'm']]
cages = {
'a': AddsTo(11),
'b': DividesTo(2),
'c': MultipliesTo(20),
'd': MultipliesTo(6),
'e': SubtractsTo(3),
'f': DividesTo(3),
'g': MultipliesTo(240),
'h': MultipliesTo(6),
'i': MultipliesTo(6),
'j': AddsTo(7),
'k': MultipliesTo(30),
'l': MultipliesTo(6),
'm': AddsTo(9),
'n': AddsTo(8),
'o': DividesTo(2),
}
answer = solve_board(board, cages)
assert answer == [[5, 6, 3, 4, 1, 2],
[6, 1, 4, 5, 2, 3],
[4, 5, 2, 3, 6, 1],
[3, 4, 1, 2, 5, 6],
[2, 3, 6, 1, 4, 5],
[1, 2, 5, 6, 3, 4]]