/
puzzles_common.py
50 lines (40 loc) · 1.37 KB
/
puzzles_common.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
import itertools
from z3 import And, Distinct
from collections import defaultdict
#transpose a square matrix
transpose = lambda m: list(zip(*m))
def flatten(list_of_lists):
return list(itertools.chain(*list_of_lists))
# index = line, column
def inside_board(index_lc, *, height, width):
l, c = index_lc
return (0 <= l < height) and (0 <= c < width)
def gen_latin_square_constraints(matrix, order):
assert len(matrix) == order
assert len(transpose(matrix)) == order
numbers = itertools.chain(*matrix)
range_c = [ And(n >= 1, n <= order) for n in numbers ]
row_c = [ Distinct(row) for row in matrix ]
col_c = [ Distinct(row) for row in transpose(matrix) ]
return range_c + row_c + col_c
# We group by content. all cells containing 0 for example
def get_same_block_indices(matrix):
res = defaultdict(list)
for l, row in enumerate(matrix):
for c, val in enumerate(row):
res[val].append((l, c))
return res
# serve first the rows then columns. (whatever their length be)
def rows_and_cols(matrix):
# serve rows
yield from matrix
# serve columns. same as yield from tranpose(matrix)
yield from zip(*matrix)
# orthogonal neighbours
def ortho_neighbours(index_):
up, down, left, right = (1, 1, 1, 1)
l, c = index_
return [(l-up, c),
(l+down, c),
(l, c-left),
(l, c+right)]