LAC-Z3-Introduction/Lecture3/colours.py

from itertools import combinations
from z3 import *

solver = Solver()

Colours = Datatype("Colours")
Colours.declare("RED")
Colours.declare("BLUE")
Colours.declare("GREEN")
Colours.declare("YELLOW")

Colour = Colours.create()
f    = Function('f', IntSort(), Colour)

variables = list()
for i in range(0,5):
    variables.append(Int(i))
    solver.add(0 <= variables[-1])
    solver.add(variables[-1] <= 5)

solver.add(Distinct(variables))

for combi in combinations(variables,2):
    solver.add(Implies(Abs(combi[0] - combi[1]) == 1, f(combi[0]) != f(combi[1])))

result = solver.check()
if result == sat:
    print(solver.model())