Power programming

Yesterday I read called Power Programming, in which the author gives solutions for this Google Code Jam problem in Python, Perl, Arc and C++. I figured I’d have a go at providing a solution in Haskell to see how it stacks up.

Since converting the given tree to a tuple won’t work in Haskell (it might with Data.Dynamic, but that’s not exactly standard practice), we’ll have to settle for writing a parser. Fortunately, Parsec makes this really easy.

First, we need some imports.

import Control.Applicative ((<*>), (<\$>))
import Text.Parsec
import Text.Parsec.Token
import Text.Parsec.Language
import Text.Printf

We’ll also have to define the structure of the decision tree.

data Tree = Node Double String Tree Tree | Leaf Double

Since we need a parser for the decision trees anyway, we might as well write a parser for the complete test case input. Because the tokens we’ll be getting are pretty much the same as in most programming languages, we can just use the existing haskell tokenizer to parse the tree.

input    = count' testCase
testCase = (,) <\$> (natural h >> tree) <*> count' animal
tree     = parens h \$ try node <|> leaf
node     = Node <\$> float h <*> identifier h <*> tree <*> tree
leaf     = Leaf <\$> float h
animal   = identifier h >> count' (identifier h)
count' p = flip count p =<< (fromIntegral <\$> natural h)

Once we have the tree and the animals, calculating the cuteness of one of them is just a matter of taking the correct branches and multiplying all the values.

cute (Leaf x) _        = x
cute (Node x f l r) fs = x * cute (if elem f fs then l else r) fs

Showing the result just requires a bit of printf use.

output = mapM_ (\(i, (t, as)) -> printf "Case #%d:\n" i >>
mapM_ (printf "%1.7f\n" . cute t) as) . zip [1::Int ..]

And finally a function that combines the required steps.

solve = either print output . parse input ""

A test to see if everything works correctly (for the sake of brevity, we read the input from a file, but using getContents to read from the console or a plain string literal will work as well):

main = solve =<< readFile "input.txt"

Not bad I reckon. Obivously, it’s not quite as brief as the Perl solution, but at least to me it’s a whole lot more readable. It’s roughly the same size as the Arc solution, which seems about right to me.