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))

"""
|cell 1|cell 2|cell 3|cell 4|cell 5|

"""

# neighbouring cells needs to be coloured differently

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