Solution to Word Numbers (Haskell)

Worddy.hs

{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NamedFieldPuns #-}

-- Chris Yuen <chris@kizzx2.com>
-- ITA Software's "Word Numbers" puzzle (http://www.itasoftware.com/careers/puzzle_archive.html)

import Data.Function
import Data.List

-- | For this problem we'll trie like structure. An "umbrella" value is 
-- the total value that this node represents.
data Node = Node
    { nodeLabel          :: String
    , nodeValue          :: Integer
    , nodeUmbrellaCount  :: Integer
    , nodeUmbrellaLength :: Integer
    , nodeUmbrellaSum    :: Integer
    , nodeConts          :: [Node]
    } deriving (Eq)

instance Ord Node where
    compare = compare `on` nodeLabel

-- | For values under one thousand, we brute force it. This simplifies things
-- as we can divide parts into 3-digit chunks evenly.
wordify :: Int -> String
wordify n
    | n < 20    = as !! n
    | n < 100   = bs !! (n `quot` 10) ++ wordify (n `rem` 10)
    | n < 1000  = as !! (n `quot` 100) ++ "hundred" ++ wordify (n `rem` 100)
    | otherwise = undefined
    where
        as = "" : words ("one two three four five six seven " ++
            "eight nine ten eleven twelve thirteen fourteen " ++
            "fifteen sixteen seventeen eighteen nineteen")
        bs = "" : "" : words ("twenty thirty forty fifty sixty " ++
            "seventy eighty ninety")

length' :: String -> Integer
length' = fromIntegral . length

numbers :: [String]
numbers = map wordify [1..999]

ones, thousands, millions :: [Node]
ones      = sort $ mkNodes [] 0 "" [1..999]
thousands = sort $ mkNodes [ones] 999 "thousand" [1000,2000..999000]
millions  = sort $ mkNodes [ones, thousands] 999999 "million" [1000000,2000000..999000000]

mkNodes :: [[Node]] -> Integer -> [Char] -> [Integer] -> [Node]
mkNodes conts umbrellaCount suffix values = zipWith go (map (++ suffix) numbers) values
    where
        sumOfSumsOfConts    = sum . map (sum . map nodeUmbrellaSum) $ conts
        sumOfLengthsOfConts = sum . map (sum . map nodeUmbrellaLength) $ conts

        addSum    = mkRecursiveAdd (\n i -> n { nodeUmbrellaSum = nodeUmbrellaSum n + i })
        addLength = mkRecursiveAdd (\n i -> n { nodeUmbrellaLength = nodeUmbrellaLength n + i })

        go label value = case conts of 
            [] -> Node label value 1 (length' label) value []
            _  -> Node label value countWithSelf totalLength totalSum (sort conts')
                where
                    labelLength   = length' label
                    countWithSelf = umbrellaCount + 1
                    totalLength   = labelLength * countWithSelf + sumOfLengthsOfConts
                    totalSum      = value * countWithSelf + sumOfSumsOfConts
                    conts'        = concatMap (addSum value . addLength labelLength) conts

-- | This is used to create a function that recursively adds 
-- a value to a field in a node
--
-- 'plus' is a function that adds an integer value to the desired field of the node
mkRecursiveAdd :: (Node -> Integer -> Node) -> (Integer -> [Node] -> [Node])
mkRecursiveAdd plus n xs = map go xs
    where
        go = if (null . nodeConts . head $ xs) then go1 else go2
        go1 x@(Node { .. }) = x `plus` n
        go2 x@(Node { .. }) =
            (plus x $ n * nodeUmbrellaCount) { nodeConts = mkRecursiveAdd plus n nodeConts }

solve :: Integer -> [Node] -> ([Char], Integer)
solve n = go 0 0 "" 0
    where
        go accLength accSum prefix prefixValue ((Node {..}):xs)
            | accLength + nodeUmbrellaLength >= n =
                if accLength + labelLength >= n
                    then (combinedLabel, accSum + combinedValue)
                    else -- Drill into
                        let accLength' = accLength + labelLength
                            accSum' = accSum + combinedValue
                        in go accLength' accSum' combinedLabel combinedValue nodeConts
            | otherwise = -- Skip
                go (accLength + nodeUmbrellaLength) (accSum + nodeUmbrellaSum) prefix prefixValue xs

            where
                combinedValue = prefixValue + nodeValue
                combinedLabel = prefix ++ nodeLabel
                labelLength   = length' combinedLabel

        go _ _ _ _ [] = error "Solution not found"

main :: IO ()
main = print . solve 51000000000 $ everything
    where everything = sort (ones ++ thousands ++ millions)