Kgtkr's Blog

ABC167 C - Skill UpをHaskellで解く

2020/05/11
proconHaskell

問題

https://atcoder.jp/contests/abc167/tasks/abc167_c

解法

解法としては各本を買う/買わないで全探索するとO(M2^N)で解けます。
C_n A_n_mは型としては[(Int, [Int])]になります。各要素の(Int, [Int])は単位元を(0, 要素が0の無限リスト)、演算をfstは足す、sndはzipして各要素を足すとするとモノイドになりそうですし問題を解くのに使えそうです。このようになる(Int, [Int])のnewtypeは(Sum Int, Ap ZipList (Sum Int))になります。ApApplicative fMonoid aからMonoid (f a)を作るnewtypeです。ZipListpurerepeat xliftA2zipして関数適用というApplicative実装を持つListのnewtypeなのでApと組み合わせると欲しいMonoidが得られます。
ところでnewtypeの変換にはcoerceという関数が便利です。安全に変換できる2つの型同士をオーバーヘッドなしで変換してくれる便利関数です。ある型からA型に変換したい時はTypeApplications拡張と組み合わせてid @A . coerceと書くと型注釈が楽です。
次に非決定性計算です。Haskellのリストモナドは非決定性計算なので全探索に便利です。

do
  x <- [1, 2, 3]
  y <- [10, 20, 30]
  pure $ x + y

とすると、[1, 2, 3][10, 20, 30]を足した時の全パターンのリストを得ることができます。
今回は買う/買わないの全列挙で、買わない場合はさっきのモノイドの単位元になるので疑似コードは以下のようになりそうです。

do
  x_1 <- [x_1, mempty]
  x_2 <- [x_2, mempty]
  ︙
  x_n <- [x_n, mempty]

  pure $ x_1 <> x_2 <> … <> x_n

こうすることで、全購入パターンの金額と各アルゴリズムの理解度を得ることができます。
これはfmap mconcat . traverse (: [mempty])で実現できます。(: [mempty])\x -> [x, mempty]です。traverseは与えられた関数でfmapしたあとにsequenceする関数です。sequenceはここではリストで与えられたモナドを「実行」して結果のリストを得る関数を考えることができます。こうすることで全購入パターンを得ることができるので、各購入パターンで金額と各アルゴリズムの理解度を得るためにfmap mconcatしています。

この後はfilter (all (>= x) . snd)で各アルゴリズムの理解度がx以上の物のみをフィルタリングし(入力xは1以上で、各アルゴリズムの理解度が無限リストの時、つまり全て変わらないパターンの時全要素は0になるので無限ループにはなりません)、map fstで金額のみを取り出してfromMaybe (-1) . foldl1May minで金額の最小値を得る、存在しなければ-1を返すだけです。
foldl1Maysafeパッケージにあるやつの特殊版でないので自分で定義しています。

全コード

{-# LANGUAGE TypeApplications #-}

import Data.Maybe
import Data.List
import Control.Monad
import Data.Bifunctor
import Control.Applicative
import Data.Monoid
import Data.Coerce

main :: IO ()
main = do
    [n, m, x] <- fmap (read @Int) . words <$> getLine
    list <- replicateM n $ fromJust . uncons . fmap (read @Int) . words <$> getLine
    print $ solve x list

solve :: Int -> [(Int, [Int])] -> Int
solve x = fromMaybe (-1)
        . foldl1May min
        . map fst
        . filter (all (>= x) . snd)
        . id @[(Int, [Int])]
        . coerce
        . fmap mconcat
        . traverse (: [mempty])
        . id @[(Sum Int, Ap ZipList (Sum Int))]
        . coerce

foldl1May :: (a -> a -> a) -> [a] -> Maybe a
foldl1May f (x:xs) = Just $ foldl f x xs
foldl1May _ _ = Nothing