Submission #5409769

---

Source Code Expand

import qualified Control.Monad as M
import qualified Data.Map as Map
import qualified Data.Set as S
import qualified Data.Array as A

type IntSet = S.Set Int
type IntArr = A.Array Int

main = do
    n <- readLn :: IO Int
    ps <- map read . words <$> getLine
    print $ solve n ps

-- solve :: Int -> [Int] -> Int
-- solve n ps = length . S.fromList $ sums where
--     sums = [sum $ zipWith (*) ptn ps | ptn <- ptns]
--     ptns = M.replicateM n [0, 1]

solve :: Int -> [Int] -> Int
solve n ps = length $ evalDP (f' ps) n
f' :: [Int] -> DP Int IntSet
f' (p:[]) = dp $ \i -> return $ S.fromList [0, p]
f' (p:ps) = dp $ \i -> do
    s <- f' ps (i-1)
    return $ S.union s (S.map (+p) s)


-- DP
type Memo a b = Map.Map a b
type DP a b = a -> State (Memo a b) b

dp :: Ord a => DP a b -> DP a b
dp f x = do
  memo <- gets (Map.lookup x)
  case memo of
       Just y  -> return y
       Nothing -> do
         y <- f x
         modify (Map.insert x y)
         return y

evalDP :: DP a b -> a -> b
evalDP f x = evalState (f x) Map.empty

evalDPAll :: DP a b -> [a] -> [b]
evalDPAll f xs = evalState (sequence (map f xs)) Map.empty

-- State
newtype State s x = State { runState :: s -> (x, s) }

instance Functor (State s) where
    fmap = M.liftM

instance Applicative (State s) where
    pure x = State $ \s -> (x, s)
    (<*>) = M.ap

instance Monad (State s) where
    m >>= k = State $ \s ->
        let (x, s0) = runState m s
        in runState (k x) s0

get :: (State s) s
get = State $ \s -> (s, s)

put :: s -> (State s) ()
put s = State $ \_ -> ((), s)

state :: (s -> (a, s)) -> (State s) a
state f = do
    s <- get
    let (a, s') = f s
    put s'
    return a

modify :: (s -> s) -> (State s) ()
modify f = get >>= put . f

gets :: (s -> a) -> (State s) a
gets f = get >>= return . f

evalState = (fst .) . runState
execState = (snd .) . runState

Submission Info

Judge Result