Data.Vector を多次元に使う idx 関数

競プロ典型 90 問 81 日目 - Friendly Group（★5）

import Control.Monad
import qualified Data.ByteString.Char8 as C
import Data.List
import Data.Maybe
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

main = ints >>= \[n,k] -> replicateM n ints >>= print . sol n k

ints = unfoldr (C.readInt . C.dropWhile (==' ')) <$> C.getLine

sol n k ts = U.maximum c 
  where
  (as,bs) = liftM2 (,) (map head) (map last) ts
  (la,ra) = (minimum as,maximum as)
  (lb,rb) = (minimum bs,maximum bs)
  wd = ra-la+1
  ht = rb-lb+1
  pos (a,b) = (a-la,b-lb)
  idx x y 
    | x<wd && y<ht  = Just $ x+wd*y
    | otherwise     = Nothing
  c = U.create $ do
    c <- UM.replicate (wd*ht) (0 :: Int)
    forM_ ts $ \[a,b] -> do
      let
        (lx,ly) = pos (a,b) 
        (rx,ry) = pos (a+k+1,b+k+1)
        f m x y = maybe (return ()) (UM.unsafeModify c m) (idx x y)
      f succ lx ly
      f succ rx ry
      f pred lx ry
      f pred rx ly
    forM_ [fromJust $ idx x y | x <- [0..wd-2], y <- [0..ht-1]] $ \i ->
      UM.unsafeRead c i >>= flip (UM.unsafeModify c) (i+1) . (+)
    forM_ [fromJust $ idx x y | x <- [0..wd-1], y <- [0..ht-2]] $ \i ->
      UM.unsafeRead c i >>= flip (UM.unsafeModify c) (i+wd) . (+)
    return c