The Foursquare Theorem

This has nothing to do with the playground game, the church, or the mobile/social/local city guide that helps you make the most of where you are. (Disclosure: I work at Foursquare.)

This is about Lagrange’s four-square theorem, which states that every natural number can be expressed as the sum of four squares. For example, $$%<![CDATA[ \newcommand\m[1]{\text{ }(\text{mod }#1)} 123456789 = 2142^2 + 8673^2 + 6264^2 + 2100^2 %]]>$$

The proof given on the Wikipedia page is only an existence proof, but I was able to find a mostly constructive proof elsewhere online. I want to present an outline of the proof along with some code that carries out the construction. Here’s a preview:

*Main> foursquare 123456789
(-2142,8673,6264,-2100)

Outline of the constructive proof

To write a number $n$ as the sum of four squares:

1. If $n$ is composite, find its prime factors and write each of the factors as the sum of four squares. Then combine them pairwise using Euler’s four-square identity, which shows that the product of two numbers that are the sum of four squares can itself be written as the sum of four squares.

2. For each prime factor $p$, find numbers $a$, $b$, $c$, and $d$ such that $a^2 + b^2 + c^2 + d^2 \equiv 0 \m p$. There are 3 cases here:

   a. If $p = 2$, then $a = b = 1$ and $c = d = 0$.

   b. If $p \equiv 1 \m 4$, find $a$ such that $a^2 + 1^2 \equiv 0 \m p$, letting $b = 1$ and $c = d = 0$.

   c. If $p \equiv 3 \m 4$, find $a$ and $b$ such that $a^2 + b^2 + 1^2 \equiv 0 \m p$, letting $c = 1$ and $d = 0$.

3. Now that we have $a^2 + b^2 + c^2 + d^2 = kp$, construct $a'$, $b'$, $c'$, and $d'$ such that $a'^2 + b'^2 + c'^2 + d'^2 = mp$ where $m \lt k$. Repeat this procedure until $a'^2 + b'^2 + c'^2 + d'^2 = p$.

The whole flow of the program looks like this:

foursquare n = factor n |> map foursquare' |> foldl1 combine
  where
    foursquare' p
      | p == 2 = (1, 1, 0, 0)
      | p `mod` 4 == 1 =
        let a = undefined
        in (a, 1, 0, 0) |> reduce
      | p `mod` 4 == 3 =
        let a = undefined
            b = undefined
        in (a, b, 1, 0) |> reduce
      where
        reduce (a, b, c, d) = undefined

So we factor n, and for each prime factor, find its four-square decomposition using foursquare', then combine them pairwise. For its part, foursquare' does some magic to find initial values for our 4 numbers and shuffles them off to reduce to get the sum of their squares to be exactly equal to p instead of a multiple of p.

I also defined the operator |> as backwards function application. It makes more sense to me to push a value through a series of functions than to have to read all my expressions inside out.

Alright, now I’m going to go through piece by piece and fill in missing definitions.

The easy parts

Factoring is probably the second Haskell program you ever wrote (the first probably also beginning with “factor”). Here it is from scratch, for fun:

nats = [1..]
divides a b = b `mod` a == 0
sieve (h:t) = h : sieve (filter (not . divides h) t)
primes = sieve (tail nats)

factor n = 
  let 
    factor' 1 _ = []
    factor' n ps@(p:pt) = if p `divides` n then p : factor' (n `div` p) ps else factor' n pt
  in factor' n primes

And yes, there are more efficient ways to generate primes.

combine is a straighforward realization of Euler’s four-square identity:

combine (a, b, c, d) (a', b', c', d') = (w, x, y, z)
  where
    w = a*a' + b*b' + c*c' + d*d'
    x = a*b' - b*a' - c*d' + d*c'
    y = a*c' + b*d' - c*a' - d*b'
    z = a*d' - b*c' + c*b' - d*a'

I’m also going to need modular exponentiation, to compute $a^n \m p$:

modexp a 0 p = 1
modexp a n p =
  let a2 = modexp (a*a `mod` p) (n `div` 2) p
  in if n `mod` 2 == 0 then a2 else a2 * a `mod` p

And a helper for finding the sum of squares of 4 numbers:

sumOfSquares (a, b, c, d) =  a*a + b*b + c*c + d*d

Before we go on, let’s check that combine works.

*Main> let x = (1, 2, 3, 4)
*Main> let y = (5, 6, 7, 8)
*Main> sumOfSquares x
30
*Main> sumOfSquares y
174
*Main> sumOfSquares (combine x y)
5220
*Main> 30*174
5220

Let’s do that 100 more times.

*Main> quickCheck (\x y -> (sumOfSquares x) * (sumOfSquares y) == sumOfSquares(combine x y))
+++ OK, passed 100 tests.

I freaking love quickCheck.

Odd primes congruent to 1 mod 4

Now to fill in the missing parts of foursquare'. Here’s the first clause:

foursquare' p
  | p `mod` 4 == 1 =
    let a = undefined
    in (a, 1, 0, 0) |> reduce

Here we have to find $a$ such that $a^2 + 1 \equiv 0 \m p$. In other words, we’re looking for a square root of $-1 \m p$. By Fermat’s little theorem, if $n \lt p$ then $n^{p-1} \equiv 1 \m p$, which means that its square root, $n^{(p-1)/2}$, is congruent to $\pm 1 \m p$. If we can find one $n$ such that $n^{(p-1)/2} \equiv -1 \m p$, then its square root, $n^{(p-1)/4}$, is the $a$ we’re after. Notice this only works because $(p-1)/4$ is an integer.

Well luckily it turns out that half of the time, $n^{(p-1)/2} \equiv -1 \m p$, so if we just try numbers sequentially we’re liable to find one pretty quickly. It also turns out that the smallest such $n$ is always a prime number, so we can narrow down the search that way.

The code is below. We’re just trying prime numbers $n$ until one of them satisifies $(n^{(p-1)/4})^2 \equiv -1 \m p$.

foursquare' p
  | p `mod` 4 == 1 =
    let
      findSqrtMinus1 (n:ps) =
        let r = modexp n ((p-1) `div` 4) p
        in if r*r `mod` p == p-1 then r else findSqrtMinus1 ps
      a = findSqrtMinus1 primes
    in (a, 1, 0, 0) |> reduce

Odd primes congruent to 3 mod 4

Here’s the next clause of foursquare':

foursquare' p
  | p `mod` 4 == 3 =
    let a = undefined
        b = undefined
    in (a, b, 1, 0) |> reduce

This case is a little more difficult. We’re looking for $a$ and $b$ such that $a^2 + b^2 + 1 \equiv 0 \m p$. Equivalently, we’re looking for a $b$ that is a square root of $-1 - a^2 \m p$. We’re guaranteed that there is some $a$ such that $-1 - a^2$ has a square root $\m p$, but I’ll refer you to the the full proof for the details on why this is. For our purposes, all we need to do is try different values of $a$ sequentially and check whether they produce a square root. That’s easy to check, at least: if $p \equiv 1 \m 4$, then $x^{(p+1)/4}$ is a square root of $x \m p$, if it has one. Here’s a quick proof:

since $(x^{\frac{1}{2}})^{p-1} \equiv 1 \m p$ by Fermat’s little theorem. (This doesn’t always produce a square root of $x \m p$ since it might instead be a square root of $-x \m p$.)

The proof gets a little hand-wavy on this part, but it turns out that finding an $a$ such that $-1 - a^2$ has a square root mod $p$ is pretty easy. It points out that -2 has a square root mod $p$ for all $p \equiv 3 \m 8$, which is already half of the cases.

Anyway, here’s the code:

foursquare' p
  | p `mod` 4 == 3 =
    let
      findNumberWithSqrt (a:ns) =
        let
          x = (-1 - a*a) `mod` p
          b = modexp x ((p+1) `div` 4) p
        in if b*b `mod` p == x then (a, b) else findNumberWithSqrt ns
      (a, b) = findNumberWithSqrt [1..p]
    in (a, b, 1, 0) |> reduce

The reduce phase

Now that we have $a^2 + b^2 + c^2 + d^2 = kp$, we want to somehow construct a new set of numbers $a'$, $b'$, $c'$, and $d'$ such that $a'^2 + b'^2 + c'^2 + d'^2 = mp$ where $m \lt k$. Then we can repeat the procedure until $a'^2 + b'^2 + c'^2 + d'^2 = p$.

If $k$ is even, then $a^2 + b^2 + c^2 + d^2$ is even, meaning 0, 2 or all 4 of the numbers are even. Rearranging them according to their remainder mod 2, we can get it so that $a \pm b$ and $c \pm d$ are both even. Then we can divide through by 2 and get

and we have the reduction we need. In code:

reduce abcd@(a, b, c, d) =
  let k = (sumOfSquares abcd) `div` p
  in case k `mod` 2 of
    0 ->
      let [a', b', c', d'] = sortWith (`mod` 2) [a, b, c, d]
      in (a'+b', a'-b', c'+d', c'-d') |> map4 (`div` 2)
    1 -> undefined

making use of the convenience functions

sortWith f = sortBy (\a b -> compare (f a) (f b))
map4 f (a, b, c, d) = (f a, f b, f c, f d)

If $k$ is odd, we can replace each of the 4 numbers by their “absolute least residue” mod $k$. All this means is, if $a \gt \frac{k}{2}$, we replace it with $a - k$, which reduces its absolute value without changing its meaning mod $k$. So let $a'$ be the absolute least residue of $a \m k$ and so on. Since all of them are smaller than $\frac{k}{2}$, their squares are less than $\frac{k^2}{4}$ and the sum of their squares is less than $k^2$. The sum of their squares is still $0 \m k$, so the sum must be equal to $mk$ for some $m \lt k$.

Now we can use combine to construct the product

The combine function is written in such a way that each of the terms $w$, $x$, $y$, and $z$ will be divisible by $k$. This is easy to see by inspecting each term and remembering that $a \equiv a' \m k$, etc. So we can divide out $k$ from each of the terms to get a product equal to $mp$, and we have the reduction we need.

Here’s the complete code for reduce:

reduce abcd@(a, b, c, d)
  | sumOfSquares abcd == p = abcd
  | otherwise =
    let k = (sumOfSquares abcd) `div` p
    in case k `mod` 2 of
      0 ->
        let [a', b', c', d'] = sortWith (`mod` 2) [a, b, c, d]
        in (a'+b', a'-b', c'+d', c'-d') |> map4 (`div` 2) |> reduce
      1 ->
        abcd |> map4 (absoluteLeastResidue k) |> combine abcd |> map4 (`div` k) |> reduce
  where
    absoluteLeastResidue k m = if a <= k `div` 2 then a else a - k
      where a = m `mod` k

And that completes the construction.

Let’s try it out:

*Main> foursquare 2013
(10,12,40,13)
*Main> foursquare 20131114
(4455,533,0,0)
*Main> foursquare 1729
(32,-22,-11,10)
*Main> quickCheck (\n -> n >= 0 ==> sumOfSquares (foursquare n) == n)
+++ OK, passed 100 tests.

Fun times.

So yeah, this is pretty useless, but it was fun to see it actually work. The code is available in this gist if you want to play around with it.