# Jane Street Puzzle, April 2022: Almost Magic

In this puzzle from Jane Street, we must fill 4 almost magic squares using distinct numbers with minimum sum.

For the purposes of this puzzle, define a magic square to be a 3-by-3 grid of positive integers for which the rows, columns, and two diagonals all have the same sum.
An almost magic square is, well, almost a magic square. It differs from a magic square in that the 8 sums may differ from each other by at most 1.
For this puzzle, place distinct positive integers into the empty grid above such that each of four bold-outlined 3-by-3 regions is an almost magic square. Your goal is to do so in a way that minimizes the overall sum of the integers you use.

This problem can be solved by integer linear programming, using, for example, mip-python.
For each square $k$ (numbered from top to bottom) and for each $i, j$ coordinates, we define an integer variable $x_{i, j, k}$. Since there are overlaps between the squares, we equalize the variables in the same region:

In [1]:
from mip import *
import itertools

m, x = Model(), {}
for k, i, j in itertools.product(range(4), range(3), range(3)):
if (i, j, k) in [(0, 0, 1), (1, 0, 1)]:
x[i, j, k] = x[i + 1, 2, 0]
elif (i, j, k) in [(0, 1, 2), (0, 2, 2)]:
x[i, j, k] = x[2, j - 1, 0]
elif (i, j, k) in [(0, 0, 3), (1, 0, 3)]:
x[i, j, k] = x[i + 1, 2, 2]
elif (i, j, k) in [(0, 1, 3), (0, 2, 3)]:
x[i, j, k] = x[2, j - 1, 1]
else:
x[i, j, k] = m.add_var(var_type=mip.INTEGER, name=f"x{i}{j}{k}")
m += x[i, j, k] >= 1


Then, we make sure that two variables $x, y$ refering to distinct cases are different. To do this, we introduce a binary variable $e$ with the following constraints: $$x \geq y + 1 - M\times e$$ $$y \geq x + 1 - M\times (1 - e)$$ With $M$ a big constant (considered as infinity). Depending on the value of $e$, these constraints reduce to:

• if $e = 0$: $x \geq y + 1$ (the other constraint $x \geq -\infty$ being useless)
• if $e = 1$: $y \geq x + 1$

Therefore those equations indeed ensure that $x, y$ are distinct.

In [2]:
e = {}

for i1, j1, k1 in itertools.product(range(3), range(3), range(4)):
for i2, j2, k2 in itertools.product(range(3), range(3), range(k1 + 1)):
if (i1, j1, k1) != (i2, j2, k2) and x[i1, j1, k1].name != x[i2, j2, k2].name:
e[i1, j1, k1, i2, j2, k2] = m.add_var(var_type=mip.BINARY, name=f"e{i1}{j1}{k1}{i2}{j2}{k2}")
m += x[i1, j1, k1] >= 1 + x[i2, j2, k2] - e[i1, j1, k1, i2, j2, k2]*150
m += x[i2, j2, k2] >= 1 + x[i1, j1, k1] - (1 - e[i1, j1, k1, i2, j2, k2])*150


The "almost magic" property can be enforced by requiring every line, column and diagonal to have the same sum as the first line, $\pm$ 1:

In [3]:
def pm(c, l, e): # add constraint c = l +/- 1
global m
m += c <= l + e
m += c >= l - (1 - e)

for k in range(4):
l = x[0, 0, k] + x[0, 1, k] + x[0, 2, k] # first line
for i in range(3):
if i != 0:
pm(x[i, 0, k] + x[i, 1, k] + x[i, 2, k], l, e[k])
pm(x[0, i, k] + x[1, i, k] + x[2, i, k], l, e[k])
pm(x[0, 0, k] + x[1, 1, k] + x[2, 2, k], l, e[k])
pm(x[0, 2, k] + x[1, 1, k] + x[2, 0, k], l, e[k])


Let's check the size of the model:

In [4]:
print(f"Number of variables: {len(m.vars)}")
print(f"Number of constraints: {len(m.constrs)}")

Number of variables: 798
Number of constraints: 1616


Finally, let's optimize the sum:

In [5]:
m.objective = minimize(xsum(x[i, j, k] for i, j, k in itertools.product(range(3), range(3), range(4))))
m.verbose = 0
m.optimize(max_seconds=1200)

Out[5]:
<OptimizationStatus.FEASIBLE: 3>

Here's the computed solution, for each square:

In [7]:
for k in range(4):
for i in range(3):
for j in range(3):
print(f"{x[i, j, k].x}", end=" ")
print()
print()

print(f"Sum of integers: {m.objective_value}")

14 17 3
1 11 22
19 5 9

22 40 8
9 23 37
39 7 24

10 19 5
6 12 16
18 4 13

16 39 7
13 21 29
34 2 26

Sum of integers: 600.0


It turns out that this is an optimal solution, even though mip could not prove it in a reasonable time.

Tags:

Updated: