2008 Round 1B

Crop Triangles

The center is a grid point if and only if both coordinates are divisible by 3. Thus we reduce all coordinates modulo 3, so each tree can be at one of 9 locations. We allow multiple trees at the same location.

Given three trees whose center is a grid point, if two of them are at the same location, the third must be as well. Thus:

  1. Choose 3 trees at the same location. If there are n trees in a given location, there are choose n 3 ways of doing this.

  2. Choose 3 trees at 3 distinct locations that sum to 0 mod 3. The number of ways of doing this is simply the product of the number of trees at each location.

This is enough to solve the problem. The number 9 is low enough that we can ineffiicently enumerate all triples of distinct locations with subsequences.

import Data.List
import qualified Data.Map as M
import Jam

main = jam $ do
  [n, a, b, c, d, x0, y0, m] <- getints
  let
    xs = take n $ unfoldr (\x -> Just (x, (a*x + b) `mod` m)) x0
    ys = take n $ unfoldr (\y -> Just (y, (c*y + d) `mod` m)) y0
    ps = M.fromListWith (+) $ zip
      (zip (map (`mod` 3) xs) (map (`mod` 3) ys)) (repeat 1)
    fwd r = M.findWithDefault 0 r ps
  return $ show $ sum ((\n -> n * (n - 1) * (n - 2) `div` 6) <$> ps) +
    sum [product $ map fwd ts | ts <- subsequences $ M.keys ps,
      length ts == 3, is3 fst ts, is3 snd ts]

is3 f ts = sum (map f ts) `mod` 3 == 0

We gain a little efficiency by observing the x-coordinates of three trees sum to 0 modulo 3 if and only if they are all the same or all different, and similarly for the y-coordinates, but there’s no need for this optimization in a contest.

Number Sets

The problem description strongly suggests something like the Sieve of Erastothenes should work. The largest common prime factor we need to consider is at most b - a + 1 (at most one multiple of any larger prime can lie in the interval), which suggests we should run some algorithm starting from the smallest prime not less than p, then work our way up to the largest prime not greater than b - a + 1.

This approach turns out to work, but I spent a long time on this problem for two reasons. Firstly, my algorithm was wrong, in the worst possible way. I reasoned thus:

  1. Start with the number of sets equal to b - a + 1. As a first approximation, each number lies in its own singleton set.

  2. For a prime p, we mark off its multiples in [a..b] if they are unmarked. Suppose we mark n numbers. If alaredy marked, we set a boolean meaning that at least one collision occurred.

  3. If the boolean is false, that is, no collisions occurred, then subtract n - 1 from the number of sets because n - 1 of these numbers actually lie in the same set as the first.

  4. Otherwise if at least one of the numbers were already marked, then subtract n because each of these numbers lies in at least one other set.

  5. If we’ve processed all the primes within range, then output the number of sets we’ve computed. Otherwise take a prime we have yet to do, and go back to step 2.

It turns out this algorithm works for the small input file, which is why its fatal flaw is especially pernicious. Fooled by its early success, I was dumbstruck when it failed on the large input file.

import Data.Array
import Data.List
import Data.Numbers.Primes
import Jam

cdiv a n = (a + n - 1) `div` n

flawedMain = jam $ do
  [a, b, p] <- getints
  let
    ps = takeWhile (<= b - a + 1) $ dropWhile (< p) primes
    tmp = listArray (a, b) $ repeat False
    f (acc, tmp) p = let
      start = p * cdiv a p
      g (fresh, collide, tmp) i = if tmp!i then (fresh, True, tmp)
        else (fresh + 1, collide, tmp // [(i, True)])
      (f, c, tmp') = foldl' g (0, False, tmp) [start, start + p..b]
      in (acc + if c then f else f - 1, tmp')

  return $ show $ b - a + 1 - fst (foldl' f (0, tmp) ps)

What is the flaw? We’ll leave it as an exercise, but here’s a hint: suppose the multiples of a prime p1 and the multiples of another prime p2 are mutually exclusive. Furthermore, suppose when marking off multiples of a prime p3, we find it collides with two different numbers.

Now, if one collision is a multiple of p1 and the other a multiple of p2, then we have one big set containing multiples of each of three primes. On the other hand, if both collisions are multiples of p1 then we have two sets: one containing multiples of p1 and p3, and the other containing multiples of p2. Thus one case has one more set than the other. The above algorithm gives the same answer in both cases, so it must be wrong.

To fix our algorithm, we must abandon our attempted shortcut and admit we need to track which primes collide with which other primes. Instead of marking each number with a boolean, we mark each number with the first prime factor p we find; later, when we find a different prime q is also a factor, we remember that the multiples of p and q lie in the same set.

The Galler-Fisher disjoint set data structure efficiently solves this problem. Each prime starts in its own singleton set and we perform disjoint set unions for each collision we find.

Our answer is then the number of disjoint sets, plus the number of unmarked numbers, with one complication: we check a prime divides at least two numbers within the interval before marking its multiples.

The other reason I had trouble with this problem was the opposite of premature optimization (overdue optimization?). Referential transparency and succinct expressive notation have a cost. Array updates, for example, are cheap in conventional languages, but relatively expensive in pure settings because nothing is mutable.

The cost is prohibitive in this problem. Although Haskell readily handles mutable state, I stubbornly hoped I could stick to clean pure code, and tried to replace the marking of numbers with a routine that counted the multiples of a given factor within the interval using a couple of divisions and used the inclusion-exclusion principle and a recursion to consider products of [2..] primes at a time.

Again, I had to abandon an attempted shortcut. I resorted to using the ST monad, which gives us fast array updates and in general, the equivalent of variables in conventional languages. The price is uglier code: we lose referential transparency and we must be more explicit, which means bugs are more likely.

import Data.STRef
import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import qualified Data.IntDisjointSet as DJ
import Data.List
import Data.Numbers.Primes
import Jam

cdiv a n = (a + n - 1) `div` n

main = jam $ do
  [a, b, p] <- getints
  let
    lim = b - a + 1
    ps = takeWhile (<= lim) $ dropWhile (< p) primes
  return $ show $ runST $ do
    tmp <- newListArray (a, b) $ repeat 0 :: ST s (STUArray s Int Int)
    dj <- newSTRef DJ.empty
    singles <- newSTRef lim
    sequence_ [do
      t <- readArray tmp q
      if t /= 0 then
        modifySTRef' dj $ DJ.union p t . DJ.insert p
      else do
        modifySTRef' singles (+ (-1))
        modifySTRef' dj $ DJ.insert p
        writeArray tmp q p
      | p <- ps, let s = p * cdiv a p, s + p <= b, q <- [s, s + p..b]]
    _dj <- readSTRef dj
    count <- readSTRef singles
    return $ count + DJ.disjointSetSize _dj

Mousetrap

One obvious strategy is to simulate the game on a list [1..k]: on the ith move, we remove the corresponding element and place it in a map, giving it the value i. We can then read off the desired cards through fast map lookups. For large inputs, traversing the list to build this map is far too slow, but this solution is useful for the small input and also for testing.

import Data.Array
import Data.List
import Jam

slowMain = jam $ do
  [k] <- getints
  (_:ds) <- getints
  let
    f _ [] = []
    f n a = b:f (n+1) (bs++as) where
      (as, (b:bs)) = splitAt ((n - 1) `mod` (k - n + 1)) a
    xs = array (1, k) $ zip (f 1 [1..k]) [1..k]
  return $ unwords $ map (show . (xs!)) ds

We can speed up our program a little with arrays and then perhaps with the ST monad, but ultimately we’ll need a smarter approach to tackle the large inputs.

The problem description sounds like the Josephus Problem, and the same dynamic programming trick solves both problems.

Briefly, we simulate the first move, and then recurse to solve the rest after adjusting the indices with simple modular arithmetic.

import Jam

main = jam $ do
  [k] <- getints
  (_:ds) <- getints
  let
    f n k d = case mod (d - n) k of
      0 -> n
      d1 -> f (n + 1) (k - 1) d1
  return $ unwords $ map (show . f 1 k) ds

Ben Lynn blynn@cs.stanford.edu 💡