Decision Tree

Parsec is a Haskell library well-suited for grammars like those in this problem:

import Jam
import Data.List
import Data.Tree
import Text.ParserCombinators.Parsec
import Text.Printf

main = jam $ do
  [n] <- getints
  ss <- getsn n
  [m] <- getints
  qss <- map (tail . tail . words) <$> getsn m
  let
    Right root = parse dtree "" $ concat ss
    g (Node (q, _) [])     pr qs = pr * q
    g (Node (q, s) [t, f]) pr qs = g (if elem s qs then t else f) (pr * q) qs
  pure $ concat $ ("\n" ++) . printf "%.7f" . g root 1 <$> qss

dtree :: Parser (Tree (Double, String))
dtree = do
  ws
  char '('
  ws
  wt <- many1 (digit <|> char '.')
  ws
  (id, kids) <- option ("", []) subtree
  char ')'
  ws
  return $ Node (read wt, id) kids

subtree = (,) <$> many1 letter <*> sequence [dtree, dtree]

ws = many (oneOf " \n")

The Next Number

A generalization of the problem of finding the next permutation lexicographic order.

We look for the rightmost pair of adjacent digits that appear in increasing order.

If no such pair exists, then the digits are in decreasing order, so we are forced to add a zero digit. Thus the next number can be constructed by starting with the minimum nonzero digit of the number, followed by all the zeros, including the zero we just added, then the remainder of the digits in increasing order.

We take tiny shortcuts: since the digits are in decreasing order the minimum nonzero digit is the rightmost nonzero digit, and we can simply reverse the other digits to sort.

Otherwise let (a, b) be the rightmost pair of adjacent digits such that a < b. Then we find the smallest digit a' exceeding a to its right, and swap them. Lastly we reverse the digits to the right of a', which puts them in increasing order.

import Jam
import Data.List

main = jam $ do
  cs <- gets
  let spls = drop 2 $ reverse $ zip (inits cs) (tails cs) of
  pure $ case find (\(_, a:b:_) -> a < b) spls of
    Nothing         -> b:'0':as ++ bs
      where (as, b:bs) = span (== '0') $ reverse cs
    Just (xs, y:ys) -> xs ++ b:as ++ y:bs
      where (as, b:bs) = span (<= y)   $ reverse ys

Square Math

Brute force works on the small input. We iteratively try all paths of a certain length to a given cell. If two paths to the same destination have the same sum, we discard the longer path, or the lexicographically greater path in case of a tie.

At each step, we see if any paths found so far sum to the first query. If so, we find the best path, then remove the first query. We repeat until there are no queries left.

import Jam
import Control.Arrow
import Data.Array
import Data.Char
import Data.List
import qualified Data.Map as M
import Data.Maybe

add (a, b) (c, d) = (a + c, b + d)
dirs = [(-1, 0), (1, 0), (0, -1), (0, 1)]

toNum '+' d = digitToInt d
toNum '-' d = -digitToInt d

main = jam $ do
  [w, _] <- getints
  a <- listArray ((1, 1), (w, w)) . concat <$> getsn w
  qs <- getints
  let
    seed = M.fromList $ second f <$> filter (isDigit . snd) (assocs a)
      where f d = M.singleton (digitToInt d) [d]
    go (p, m) = [(p2, (n + dn, s ++ [c1, c2])) | p1 <- nbrs p, p2 <- nbrs p1,
      (n, s) <- M.assocs m,
      let [c1, c2] = map (a!) [p1,p2], let dn = toNum c1 c2]
    nbrs p = filter (inRange $ bounds a) $ add p <$> dirs
    best s t | length s < length t = s
             | length s > length t = t
             | otherwise           = min s t
    ins m (p, (n, s)) = M.insert p (M.insertWith best n s $ m M.! p) m
    search []     acc _ = acc
    search qs@(q:rest) acc m
      | not $ null ans = search rest (acc ++ [foldl1' best ans]) m
      | otherwise      = search qs    acc                        m'
      where
        ans = catMaybes $ M.lookup q <$> M.elems m
        m' = foldl' ins m $ concatMap go $ M.assocs m
  pure $ concatMap ("\n" ++) $ search qs [] seed

For the large input, a few optimizations were enough to bring the running time to about 6 or 7 minutes on my laptop:

  • We use strict maps.

  • We use the specialized IntMap for one of the maps.

  • In the nth step, we only consider paths of length n.

  • In each step, we try all pending queries, and remove those that match.

  • We prepend e.g. "1+" to "2-3" instead of appending "-3" to "1+2", because we use String.

However, this is almost too slow for the contest. Was I supposed to find some clever trick to speed things up?

So far I’ve only found one shortcut. First, divide each digit by the greatest common divisor of all digits in the square. Next, let b be the biggest digit that appears next to a + sign.

Suppose we find b consecutive positive integers whose best expressions all contain +b. Then we can show that the best expressions for all larger integers are the same except we replace one copy of +b with the appropriate number of copies of +b.

I’m skeptical we’re supposed to exploit this, because the largest possible query is 250.

import Jam
import Control.Arrow
import Data.Array
import Data.Bool
import Data.Char
import Data.List
import qualified Data.Map.Strict as M
import qualified Data.IntMap.Strict as I
import Data.Maybe

add (a, b) (c, d) = (a + c, b + d)
dirs = [(-1, 0), (1, 0), (0, -1), (0, 1)]

main = jam $ do
  [w, _] <- getints
  a <- listArray ((1, 1), (w, w)) . concat <$> getsn w
  qs <- getints
  let
    ds = [(p, nub [(p2, (digitToInt d - bool 0 (2 * digitToInt (a!p2))
      (a!p1 == '-'), [d, a!p1])) | p1 <- nbrs p, p2 <- nbrs p1,
      d /= '0' || p2 /= p]) | (p, d) <- assocs a, isDigit d]
    seed = M.fromList [(p, I.singleton (digitToInt d) [d]) |
      (p, d) <- assocs a, isDigit d]
    nbrs p = filter (inRange $ bounds a) $ add p <$> dirs
    next m = M.fromList $ map (second $ merge m) ds
    merge m deltas = foldl1' (I.unionWith min) $ map (bump m) deltas
    bump m (p, (dn, ds)) = (ds ++) <$> I.mapKeysMonotonic (dn +) (m M.! p)

    search qs acc m
      | null qs   = acc
      | otherwise = search todo (done ++ acc) m'
      where
        m' = next m
        done = I.toList $ I.fromListWith min $ concatMap match $ M.elems m
        todo = qs \\ (fst <$> done)
        match m = catMaybes $ foo m <$> qs
        foo m q = (,) q <$> (q `I.lookup` m)

  pure $ concatMap ("\n" ++) $ (fromJust . (`lookup` search qs [] seed)) <$> qs

Ben Lynn blynn@cs.stanford.edu 💡