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)

parser

sample = '''x18 = 1,
 x0 = 18,
 x3 = 3,
 x7 = 16,
 x11 = 15,
 x13 = 39,
 x22 = 4,
 x23 = 25,
 x20 = 23,
 x1 = 8,
 x10 = 11,
 x16 = 10,
 x25 = 33,
 x21 = 32,
 x26 = 31,
 x6 = 7,
 x2 = 2,
 x12 = 24,
 x19 = 20,
 x24 = 47,
 x9 = 14,
 x4 = 5,
 x5 = 21,
 x17 = 36,
 x8 = 6,
 x14 = 19,
 x15 = 17,
 x27 = 12'''
 
 
 
import re
t = [re.findall('\d+', v) for v in sample.split(',')]

indices =[1,2,3,4,5,6,10,11,18,7,8,9,12,13,19,20,14,15,16,17,21,22,23,24,25,26,27,28]
m = {}
for tup in t:
    var, val = tup
    var=int(var)
    val = int(val)
    m[var+1] = val
    
ans = []
for index in indices:
    print(index)
    ans.append(m[index])
print(ans)