All Base exercise from Google Code Jam 2009

AllBase.hs

module Main where

import Data.Map (Map)
import qualified Data.Map as Map
import Data.List (foldl')
import Data.Maybe (fromMaybe, fromJust)
import Control.Monad (liftM, forM_)

allBase :: String -> Int
allBase xs = foldr fn 0 $ zip [0..] (xs' ++ [1])
  where
    (symbolMap, xs') = createSymbolMap xs
    -- We have to ensure that the base is at least
    -- greater than 1 (requirement 1)
    base = case Map.size symbolMap of
             1 -> 2
             n -> n
    fn (pos, x') accum = accum + ((base ^ pos) * x')

createSymbolMap :: String -> (Map Char Int, [Int])
createSymbolMap (x:xs) = (resultMap, xs') 
  where
    initialMap = Map.insert x 1 Map.empty
    (resultMap, xs', _) = foldl' fn (initialMap, [], 0) xs
    fn (map, xs', 0) x = 
      case Map.lookup x map of
        Just x' -> (map, x':xs', 0)
        Nothing -> (Map.insert x 0 map, 0:xs', 2)
    fn (map, xs', n) x = 
      case Map.lookup x map of
        Just x' -> (map, x':xs', n)
        Nothing -> (Map.insert x n map, n:xs', n + 1)

readInput :: IO ()
readInput = do
  n <- read `liftM` getLine :: IO Int
  forM_ [1..n] readCase

readCase :: Int -> IO ()
readCase n = do
  input <- getLine
  putStrLn $ "Case #" ++ show n ++ ": " ++ show (allBase input)


main :: IO ()
main = readInput