HaskellでABC461を解く

はじめに

C問題でハマってしまい今回はC問題まででした。D問題は2次元累積和までとって提出したらTLEでした。確かにO(H^2W^2)のオーダーでは間に合わないので尺取り法のような方法で減らせそうというところまで考察できましたが時間切れとなりました。

A問題

aとdの大小比較をし条件に当てはまっているかをみます。

A問題提出

main = do
  [a,d] <- readInts
  putStrLn $ bool "No" "Yes" $ a <= d

B問題

人対斧、斧対人の表をsortして同じになるかどうかを比較します。

B問題提出

main = do
  n <- readLn :: IO Int
  as <- readInts
  bs <- readInts
  let as' = sort $ zip [1..] as
  let bs' = sort $ zip bs [1..]
  putStrLn $ bool "No" "Yes" $ as' == bs'

C問題

結局時間内になんとか間に合いましたが理由が分からずハマりました。結局同じ色で一旦グループを作ってからそれぞれの最大値を取り出す必要があるところを
グループを作らず大きい順にソートしてからグループにするような処理をしてしまってました。

C問題提出

main = do
  [n,k,m] <- readInts
  cvs <- replicateM n $ do
    [c,v] <- readInts
    return (c,v)
  -- groupにする
  let xs = sort cvs
  let xs1 = map (\g -> (maximum (snd (unzip g)), fst (head g)) ) $ groupBy (\(c1,_) (c2,_) -> c1 == c2) xs
  let ys = map (\(v,c) -> (c,v) ) $ take m $ sortBy (flip compare) xs1
  let dic = M.fromList ys
  let zs = solve dic cvs
  let ans1 = sum $ snd $ unzip ys
  let ans2 = sum $ take (k-m) $ sortBy (flip compare) $ snd $ unzip zs
  print $ ans1 + ans2

solve :: M.Map Int Int -> [(Int,Int)] -> [(Int,Int)]
solve dic [] = []
solve dic (cv@(c,v):cvs)
  | M.member c dic && dic M.! c == v = solve dic1 cvs
  | otherwise = cv : solve dic cvs
  where
    dic1 = M.delete c dic

おわりに

D問題は2次元累積和を作るところまではあっさりできたのですがそこからのひと工夫が思いつかずでした。