jane street almost magic

almost_magic.py

import z3

def add_constraints(s, inp):
    for digit in inp:
        s.add(digit > 0)
        s.add(digit < 200)

    b = [(6,10,11),(9,12,13),(15,16,17),(6,9,15),(10,12,16),(11,13,17), (6,12,17),(11,12,15)]
    a = [(1,2,3),(4,5,6),(7,8,9),(1,4,7),(2,5,8),(3,6,9), (1,5,9),(3,5,6)]
    c = [(18,7,8),(19,20,14),(21,22,23),(18,19,21),(7,20,22),(8,14,23), (18,20,23),(21,20,8)]
    d = [(14,15,16),(23,24,25),(26,27,28),(14,23,26),(15,24,27),(16,25,28), (14,24,28),(16,24,26)]
    choose =[a,b,c,d]
    
    for i in range(0, len(inp)):
        for j in range(i+1, len(inp)):
            s.add(inp[i] != inp[j])

    # add the constraints
    for chosen in choose:
        for i in range(0, len(chosen)):
            t1,t2,t3 = chosen[i]
            r2 = inp[t1-1]+inp[t2-1]+inp[t3-1]
            for j in range(i+1, len(chosen)):
                t4,t5,t6 = chosen[j]
                r1 = inp[t4-1]+inp[t5-1]+inp[t6-1]
                d = r1-r2
                s.add(z3.Or(d == -1, d == 1, d==0))

    total = 0
    for d in inp:
        total += d
        
    s.add(total >= 420)
    s.add(total < 499)

    return total

def solve(inp):
    s = z3.Optimize()
    value = add_constraints(s, inp)

    #s.minimize(value)

    assert s.check() == z3.sat
    m = s.model()
    return m.eval(value), m

inp = [z3.Int('x{}'.format(i)) for i in range(28)]
print("input", inp)

tot, out = solve(inp)
print(tot, out)