CodinGame - Virus

virus.hs

import Data.List
import Data.Ord
import System.IO
import Control.Monad

type Node = Int
type Gateway = Node
type Link = (Node, Node)
type Path = [Link]

main = do
    hSetBuffering stdout NoBuffering -- DO NOT REMOVE

    [_, linkCount, gatewayCount] <- fmap (map read . words) getLine

    links <- replicateM linkCount $ fmap (toLink . map read . words) getLine::IO [Link]

    gateways <- replicateM gatewayCount (fmap read getLine)

    loop gateways links

    where
        toLink (n1 : n2 : _) = (n1, n2)

loop gateways links = do
    ((n1, n2), remainingLinks) <- fmap (deleteLink gateways links . read) getLine
    putStrLn (show n1 ++ " " ++ show n2)
    loop gateways remainingLinks

deleteLink :: [Gateway] -> [Link] -> Node -> (Link, [Link])
deleteLink gateways links agent =
    (deletedLink, delete deletedLink links)
    where
        shortestPath = head $ sortOn length $ gatewayPaths agent links gateways
        deletedLink = last shortestPath

gatewayPaths :: Node -> [Link] -> [Gateway] -> [Path]
gatewayPaths node links gateways =
    map untilGateway $ filter (any atGateway) (paths node links)
    where
        atGateway (n1, n2) = n1 `elem` gateways || n2 `elem` gateways
        untilGateway (link : path)
            | atGateway link = [link]
            | otherwise = link : untilGateway path
            
paths :: Node -> [Link] -> [Path]
paths node links = nexts
    where
        adjacents = filter (isAdjacent node) links
        nexts =
            map return adjacents
            ++
            map (\ link -> concatMap (\ next -> link : next) (paths (otherNode link) (delete link links))) adjacents
        otherNode (n1, n2) = if n1 == node then n2 else n1
        isAdjacent node (n1, n2)
            | node == n1 || node == n2 = True
            | otherwise = False
            
-- Pasted from Data.List because it's missing in the (outdated?) version here.
-- sortOn f = map snd . sortBy (comparing fst) . map (\x -> let y = f x in y `seq` (y, x))