Solving Jane Street’s Puzzles (Day 1)

Boi Mai Quach
6 min readOct 11, 2023

--

In this series, I will in turn solve each and every Jane Street’s Puzzles, using solvers. Through each post, I break down the problem-solving process step by step, demonstrating how mathematical logic and computational tools can be used to find efficient solutions. Whether you’re new to solving puzzles or an experienced solver, this series offers insights and techniques to enhance your problem-solving skills.

Photo by Esteban López on Unsplash

In this part, I will in turn solve each and every Jane Street’s Puzzles. My code will be published on my GitHub.

✨Please give me a 👏 if you are interested in my sharing, leave a comment 💬 to share your views/thoughts. ✨

Sum of Squares

The puzzle was publised at this link.

Source: https://www.janestreet.com/puzzles/sum-of-squares-index/

Place a digit in each of the 25 spots in the below 5x5 grid, so that each 5-digit number (leading zeroes are ok) reading across and reading down is divisible by the number outside the grid, trying to maximize the sum of the 25 numbers you enter. An example of a completed grid with sum 100 is presented on the right.

Please submit your answer (along with any other comments you want to provide) as an ordered pair of your sum, and your 25 numbers, reading left to right, top to bottom.

Example submission: (100,1623552460048932486847030)

Break down the problem

This is a grid-based puzzle where you have a 5x5 grid, and the goal is to fill each cell with a digit (0–9) to create 5-digit numbers when reading across or down.

According to the example:

Numbers in the fifth row (left) and first column (right)
  • For the fifth row, we have 5 values in order: 4, 7, 0, 3, 0, creating a 5-digit number: 47,030.
  • For the first column, we have 5 values in order: 1, 5, 0, 2, 4, creating a 5-digit number: 15,024.

The second requirement is that the 5-digit numbers you form have to be divisible by the respective numbers outside the grid. Taking the example above:

  • The 5-digit number in the fifth row, 47,030, has to be divisible by 5 => satisfied.
  • The 5-digit number in the first column, 15,024, has to be divisible by 6 => satisfied.

Before solving the problem, for those who have not used z3 API (python), the next part will give them a brief knowledge about z3. Otherwise, readers can skip the z3 and directly go to the solution.

Z3 in Python

Introduction

Z3 is a high performance theorem prover developed at Microsoft Research. It is the most powerful solver I have found for such flexible systems of equations. Z3 is an excellent choice now that it is released under the MIT license.

If you are not familiar with Z3, you can start at z3-solver.

Z3 can be built using Visual Studio, a Makefile or using CMake. It provides bindings for several programming languages.

Optimisation in Z3

Here’s a basic example of solving an optimization problem using z3.optimize:

In this example, we’re trying to maximize the objective function

subject to the constraint,

The solver will find the optimal values of x and y that maximize th objective function while satisfying the given constraint. Adjust the objective function and constraints as needed for your specific optimization problem.

# Create a Z3 solver
opt = Optimize()

# Create integer variables for the optimization problem
x = Int('x')
y = Int('y')

# Define the objective function to maximize (e.g., -x - 2*y)
objective = -x - 2*y

# Add the objective function to the optimizer
opt.maximize(objective)

# Add constraints (e.g., x + y <= 10)
opt.add(x + y <= 10)

# Check for a solution
if opt.check() == sat: # return sat for satisfiable instances.
model = opt.model()
print('Optimal solution:')
print('x =', model[x].as_long())
print('y =', model[y].as_long())
print('Objective value:', model.evaluate(objective).as_long())
else:
print('No solution found.')

And the output is

Solution

Import necessary modules and packages

import numpy as np
import time
import seaborn as sns
import matplotlib.pyplot as plt
from IPython.display import Markdown
from z3 import *

Create a z3 solver

# Create a Z3 solver
opt = Optimize()

Create a 5x5 grid

# Setup up a 5x5 matrix of integer variables
X = np.array(IntVector("x",25),dtype=object).reshape((5,5))

Initialise solver

In this case, we use optimize function because the problem involves finding the maximum value of a function.

# Initialise an optimiser
s = Optimize()

Maximize the sum of the 25 numbers

As the requirements, we will maximize the sum of the 25 numbers that we enter.

# Maximise sum
s.maximize(Sum([x for _,x in np.ndenumerate(X)]))

Set up the values for all cells

The goal is to fill each cell with a digit (0–9), thus we obtain

# each cell contains a value in {0, ..., 9}
s += [And(x>=0,x<=9) for _,x in np.ndenumerate(X)]

Turn rows and columns to a 5-digit numbers

#function to turn row/col to a 5-digit number
mult = lambda x:Sum([x[i] *10**((4)-i) for i in range(5)])

For example, the number filling in the first row will be: 1, 6, 2, 3, and 5.

Numbers in the first row

In terms of the function, we obtain:

The number is

Divisibility

According to the second condition, every 5-digit number, allowing for leading zeroes, that is formed by reading horizontally or vertically must be evenly divisible by the number situated outside the grid.

Rows

The numbers outside the rows are respectively: 1, 2, 3, 4, 5

# Rows
s += [mult(X[i,:]) % (i+1)==0 for i in range(5)]

Columns

The numbers outside the rows are respectively: 6, 7, 8, 9, 10

# Cols
s += [mult(X[:,i]) % (i+6)==0 for i in range(5)]

Now, it’s time to find the solution

Processing

start = time.time()
# Solve and print
if s.check() == sat:
m = s.model()
evalu = np.vectorize(lambda x:m.evaluate(x).as_long())
r = evalu(X)
else:
print("Failed")

display(Markdown("**Solved in {:.4f} seconds. Answer is : {}**".format(time.time()-start,np.sum(r))))
fig,ax = plt.subplots(1,1,figsize=(3,3))
ax = sns.heatmap(np.zeros_like(r),annot=r,cbar=False,cmap="Greys",linecolor='k',linewidth=2,
xticklabels=[6,7,8,9,10], yticklabels=[1,2,3,4,5], annot_kws={"fontsize":18})
ax.tick_params(left=False, bottom=False,labelleft=False, labelright=True)
plt.xticks(rotation=0,fontsize =12)
plt.yticks(rotation=0,fontsize =12)
plt.tight_layout()
plt.show()
fig.savefig('sum_of_squares_result.png')

print(np.sum(r),",","".join([str(x) for _,x in np.ndenumerate(r)]))

Result

Solved in 1.0787 seconds. Answer is : 205

Public solution

The BEST SOLUTION was publised at this link.

Source: https://www.janestreet.com/puzzles/sum-of-squares-solution/

Next time, we will dive into the Hooks problem. Try to solve it by your-self first. See you next time with the detailed solution and explanation for that problem.

Note: Unless otherwise noted, all images are by the author.

Source code

https://github.com/Tayerquach/quant_puzzles

References

https://www.janestreet.com/puzzles/

https://microsoft.github.io/z3guide/docs/logic/intro/

https://github.com/gowen100/Jane-Street-Solutions

--

--

Boi Mai Quach

PhD Student at SFI Centre for Research Training in Machine Learning, Dublin City University (DCU), Ireland.