14 November 2013

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,

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

Outline of the constructive proof

To write a number as the sum of four squares:

  1. If 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 , find numbers , , , and such that . There are 3 cases here:

    a. If , then and .

    b. If , find such that , letting and .

    c. If , find and such that , letting and .

  3. Now that we have , construct , , , and such that where . Repeat this procedure until .

The whole flow of the program looks like this:

foursquare n = factor n |> map foursquare' |> foldl1 combine
    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
        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 = 
    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)
    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 :

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
*Main> sumOfSquares y
*Main> sumOfSquares (combine x y)
*Main> 30*174

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 such that . In other words, we’re looking for a square root of . By Fermat’s little theorem, if then , which means that its square root, , is congruent to . If we can find one such that , then its square root, , is the we’re after. Notice this only works because is an integer.

Well luckily it turns out that half of the time, , so if we just try numbers sequentially we’re liable to find one pretty quickly. It also turns out that the smallest such is always a prime number, so we can narrow down the search that way.

The code is below. We’re just trying prime numbers until one of them satisifies .

foursquare' p
  | p `mod` 4 == 1 =
      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 and such that . Equivalently, we’re looking for a that is a square root of . We’re guaranteed that there is some such that has a square root , 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 sequentially and check whether they produce a square root. That’s easy to check, at least: if , then is a square root of , if it has one. Here’s a quick proof:

since by Fermat’s little theorem. (This doesn’t always produce a square root of since it might instead be a square root of .)

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

Anyway, here’s the code:

foursquare' p
  | p `mod` 4 == 3 =
      findNumberWithSqrt (a:ns) =
          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 , we want to somehow construct a new set of numbers , , , and such that where . Then we can repeat the procedure until .

If is even, then 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 and 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 is odd, we can replace each of the 4 numbers by their “absolute least residue” mod . All this means is, if , we replace it with , which reduces its absolute value without changing its meaning mod . So let be the absolute least residue of and so on. Since all of them are smaller than , their squares are less than and the sum of their squares is less than . The sum of their squares is still , so the sum must be equal to for some .

Now we can use combine to construct the product

The combine function is written in such a way that each of the terms , , , and will be divisible by . This is easy to see by inspecting each term and remembering that , etc. So we can divide out from each of the terms to get a product equal to , 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
    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
*Main> foursquare 20131114
*Main> foursquare 1729
*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.

blog comments powered by Disqus