forked from 0vercl0k/z3-playground
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patheinstein_riddle_z3.py
154 lines (126 loc) · 7.86 KB
/
einstein_riddle_z3.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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# einstein_riddle_z3.py - Solve the Einstein's riddle to step into the 2% :)).
# Copyright (C) 2012 Axel "0vercl0k" Souchet - http://www.twitter.com/0vercl0k
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
from __future__ import print_function
import sys
from z3 import *
# Thx to alpounet for the idea!
def solve_einstein_stuff():
"""Solve this stuff: http://www.davar.net/MATH/PROBLEMS/EINSTEIN.HTM"""
def column(matrix, i):
"""Get the column i of matrix"""
return [matrix[j][i] for j in range(5)]
def instanciate_int_constrained(name, s):
x = Int(name)
# Each int represent an index in p[name]
s.add(x >= 0, x <= 4)
return x
p = {
'color' : ('Yellow', 'Green', 'Blue', 'Red', 'White'),
'nationality' : ('Brit', 'Dane', 'German', 'Norwegian', 'Swede'),
'beverage': ('Beer', 'Coffee', 'Milk', 'Tea', 'Water'),
'smoke' : ('Blue Master', 'Dunhill', 'Pall Mall', 'Prince', 'Blend'),
'pet' : ('Cat', 'Bird', 'Dog', 'Fish', 'Horse')
}
s = Solver()
color, nationality, beverage, smoke, pet = range(5)
houses = [[instanciate_int_constrained('%s%d' % (prop, n), s) for prop in p.keys()] for n in range(5)]
for i in range(5):
# Each column must have different values
s.add(Distinct(column(houses, i)))
# Time to feed the solver with hints
# 1. The Brit lives in a red house.
s.add(Or([And(houses[i][nationality] == p['nationality'].index('Brit'), houses[i][color] == p['color'].index('Red')) for i in range(5)]))
# 2. The Swede keeps dogs as pets.
s.add(Or([And(houses[i][nationality] == p['nationality'].index('Swede'), houses[i][pet] == p['pet'].index('Dog')) for i in range(5)]))
# 3. The Dane drinks tea.
s.add(Or([And(houses[i][nationality] == p['nationality'].index('Dane'), houses[i][beverage] == p['beverage'].index('Tea')) for i in range(5)]))
# 4. The green house is on the left of the white house (next to it).
s.add(Or([And(houses[i][color] == p['color'].index('Green'), houses[i + 1][color] == p['color'].index('White')) for i in range(4)]))
# 5. The green house owner drinks coffee.
s.add(Or([And(houses[i][color] == p['color'].index('Green'), houses[i][beverage] == p['beverage'].index('Coffee')) for i in range(5)]))
# 6. The person who smokes Pall Mall rears birds.
s.add(Or([And(houses[i][smoke] == p['smoke'].index('Pall Mall'), houses[i][pet] == p['pet'].index('Bird')) for i in range(5)]))
# 7. The owner of the yellow house smokes Dunhill.
s.add(Or([And(houses[i][color] == p['color'].index('Yellow'), houses[i][smoke] == p['smoke'].index('Dunhill')) for i in range(5)]))
# 8. The man living in the house right in the center drinks milk.
s.add(houses[2][beverage] == p['beverage'].index('Milk'))
# 9. The Norwegian lives in the first house.
s.add(houses[0][nationality] == p['nationality'].index('Norwegian'))
# 10. The man who smokes blend lives next to the one who keeps cats.
s.add(Or(
And(houses[0][smoke] == p['smoke'].index('Blend'), houses[1][pet] == p['pet'].index('Cat')),
And(houses[1][smoke] == p['smoke'].index('Blend'), Or(houses[0][pet] == p['pet'].index('Cat'), houses[2][pet] == p['pet'].index('Cat'))),
And(houses[2][smoke] == p['smoke'].index('Blend'), Or(houses[1][pet] == p['pet'].index('Cat'), houses[3][pet] == p['pet'].index('Cat'))),
And(houses[3][smoke] == p['smoke'].index('Blend'), Or(houses[2][pet] == p['pet'].index('Cat'), houses[4][pet] == p['pet'].index('Cat'))),
And(houses[4][smoke] == p['smoke'].index('Blend'), houses[3][pet] == p['pet'].index('Cat')),
))
# 11. The man who keeps horses lives next to the man who smokes Dunhill.
s.add(
Or(
And(houses[0][pet] == p['pet'].index('Horse'), houses[1][smoke] == p['smoke'].index('Dunhill')),
And(houses[1][pet] == p['pet'].index('Horse'), Or(houses[0][smoke] == p['smoke'].index('Dunhill'), houses[2][smoke] == p['smoke'].index('Dunhill'))),
And(houses[2][pet] == p['pet'].index('Horse'), Or(houses[1][smoke] == p['smoke'].index('Dunhill'), houses[3][smoke] == p['smoke'].index('Dunhill'))),
And(houses[3][pet] == p['pet'].index('Horse'), Or(houses[2][smoke] == p['smoke'].index('Dunhill'), houses[4][smoke] == p['smoke'].index('Dunhill'))),
And(houses[4][pet] == p['pet'].index('Horse'), houses[3][smoke] == p['smoke'].index('Dunhill'))
)
)
# 12. The owner who smokes Blue Master drinks beer.
s.add(Or([And(houses[i][smoke] == p['smoke'].index('Blue Master'), houses[i][beverage] == p['beverage'].index('Beer')) for i in range(5)]))
# 13. The German smokes Prince.
s.add(Or([And(houses[i][nationality] == p['nationality'].index('German'), houses[i][smoke] == p['smoke'].index('Prince')) for i in range(5)]))
# 14. The Norwegian lives next to the blue house.
s.add(
Or
(
And(houses[0][nationality] == p['nationality'].index('Norwegian'), houses[1][color] == p['color'].index('Blue')),
And(houses[1][nationality] == p['nationality'].index('Norwegian'), Or(houses[0][color] == p['color'].index('Blue'), houses[2][color] == p['color'].index('Blue'))),
And(houses[2][nationality] == p['nationality'].index('Norwegian'), Or(houses[1][color] == p['color'].index('Blue'), houses[3][color] == p['color'].index('Blue'))),
And(houses[3][nationality] == p['nationality'].index('Norwegian'), Or(houses[2][color] == p['color'].index('Blue'), houses[4][color] == p['color'].index('Blue'))),
And(houses[4][nationality] == p['nationality'].index('Norwegian'), houses[3][color] == p['color'].index('Blue'))
)
)
# 15. The man who smokes blend has a neighbor who drinks water.
s.add(
Or
(
And(houses[0][smoke] == p['smoke'].index('Blend'), houses[1][beverage] == p['beverage'].index('Water')),
And(houses[1][smoke] == p['smoke'].index('Blend'), Or(houses[0][beverage] == p['beverage'].index('Water'), houses[2][beverage] == p['beverage'].index('Water'))),
And(houses[2][smoke] == p['smoke'].index('Blend'), Or(houses[1][beverage] == p['beverage'].index('Water'), houses[3][beverage] == p['beverage'].index('Water'))),
And(houses[3][smoke] == p['smoke'].index('Blend'), Or(houses[2][beverage] == p['beverage'].index('Water'), houses[4][beverage] == p['beverage'].index('Water'))),
And(houses[4][smoke] == p['smoke'].index('Blend'), houses[3][beverage] == p['beverage'].index('Water'))
)
)
if s.check() == unsat:
raise Exception('Hmm the system does not seem solvable, I guess one of the constraint is wrong.')
m = s.model()
solution = [[m[case].as_long() for case in line] for line in houses]
for i in range(5):
print('Color: %s, Nationality: %s, Beverage: %s, Smoke: %s, Pet: %s' % (
p['color'][solution[i][color]],
p['nationality'][solution[i][nationality]],
p['beverage'][solution[i][beverage]],
p['smoke'][solution[i][smoke]],
p['pet'][solution[i][pet]])
)
def main(argc, argv):
solve_einstein_stuff()
return 1
if __name__ == '__main__':
sys.exit(main(len(sys.argv), sys.argv))