Google Code Jam/Random Route

import Data.List (unfoldr,last,groupBy,sortBy) import Data.Ord (comparing) import Control.Arrow import Data.Function import Control.Monad.State import Data.Ratio data Edge = Edge {ind:: Int , start :: City, end :: City, time :: Integer} instance Eq Edge where (==) = (==) `on` ind instance Ord Edge where compare = comparing ind type City = String type DB = Map City [Edge] type Work = (City,DB) type Route = [Edge] sgroupBy :: Ord c => (a -> c) -> [a] -> [[a]] sgroupBy f = map (map fst) . groupBy ((==) `on` snd) . sortBy (comparing snd) . map (id &&& f) bests f = head . sgroupBy f -- | make groups of edges with same ending city and take the fastest paths of each group fastest :: [Route] -> [[Route]] fastest = map (bests (sum .map time)) . sgroupBy (end .last) type Row = Map Edge (Ratio Integer) row :: [Route] -> Row row xs = foldr (flip $ foldr (flip (insertWith (+)) (1 % fromIntegral (length xs)))) empty xs -- | sum up the probability of each Edge appearing in the rows padding with fake Edges proba :: Int -> Int -> [Row] -> [String] proba k top xs = elems . flip union zeros . M.map ((\x -> longdiv (numerator x) (denominator x) k) . (* (1 % fromIntegral (length xs)))) . unionsWith (+) $ xs where zeros = fromList $ map (\n -> (Edge n undefined undefined undefined, "0." ++ replicate k '0')) [1..top] longdiv :: (Integral a,Integral b) => a -> b -> Int -> String longdiv n d k = let (ln,ld) = (fromIntegral *** fromIntegral) (n,d) (cm,t) = (id *** (> ld `div` 2)) $ divMod (ln * 10 ^ k) ld (l,cz) = (length &&& id) . show $ if t then cm + 1 else cm in if l <= k then "0." ++ replicate (k-l) '0' ++ cz else let (big,small) = splitAt (l - k) cz in big ++ "." ++ small -- | list of all routes , cutting cyles routing :: (City, Map City [Edge]) -> Map City () -> [Route] routing (x,d) seen = let s' = M.insert x () seen in if x `member` d then concatMap (\y -> [y]:(map (y:) . routing (end y,d) $ s')) . filter (not . (`member` s') . end) $ d ! x else [] parseEdge :: String -> Int -> Edge parseEdge x i = let [s,e,n] = words x in Edge i s e (read n) parseWork :: [String] -> ((Work,Int),[String]) parseWork (x:xs) = let mkDB = foldr (\x -> insertWith (++) (start x) [x]) empty [n,c] = words x (es,rest) = splitAt (read n) xs in (((c,mkDB . map (uncurry parseEdge) $ zip es [1..]),read n),rest) parseCases :: String -> [(Work,Int)] parseCases x = let (n:ts) = lines x in take (read n) . flip unfoldr ts $ \x -> if null x then Nothing else Just $ parseWork x main = do ts <- parseCases `fmap` getContents flip mapM_ (zip [1..] ts) $ \(i,(t,n)) -> do putStr $ "Case #" ++ show i ++ ": " putStrLn $ unwords $ proba 7 n $ map row $ fastest $ routing t empty