AtCoder Regular Contest 030

Submission #286413

Source codeソースコード

{-# LANGUAGE CPP #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}

import qualified Data.ByteString.Char8 as S
import Data.Monoid

main :: IO ()
main = do
  (S.lines -> [S.words -> [S.unpack -> read -> (i0 :: Int)], S.words -> [S.unpack -> read -> (i1 :: Int)]]) <- S.getContents
  S.putStrLn $ if 1 < i0 `div` i1 then "YES" else "NO"

-- This program was generated and submitted by the following program:

# if 0

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuasiQuotes #-}

-- required packages (apart from the Haskell Platform): 
--  vector-space, tagsoup, lens, interpolatedstring-perl6, curl, http-types

module Main (main, logIn) where

import Control.Applicative
import Control.Concurrent
import Control.Lens hiding (argument, op, act)
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict
import Data.AdditiveGroup
import qualified Data.ByteString.Char8 as S
import Data.Fixed
import Data.Function
import qualified Data.IntSet as IS
import Data.List
import qualified Data.Map as M
import Data.Maybe
import Data.Monoid
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.IO as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as TB
import qualified Data.Text.Lazy.Builder.Int as TB
import qualified Data.Text.Read as T
import qualified Data.Time as Time
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Debug.Trace
import Language.Haskell.TH
import qualified Network.Curl as Curl
import qualified Network.HTTP.Types.URI as HTTPTypes
import qualified Text.HTML.TagSoup as TS
import qualified Text.HTML.TagSoup.Match as TS
import Text.InterpolatedString.Perl6

main :: IO ()
main = do
  sleepUntil $ utc 2014 11 29 12 00 03.0

  solveAndSubmit "arc030"

-- | Create a UTCTime.
utc :: Integer -> Int -> Int -> Int -> Int -> Pico -> Time.UTCTime
utc year month day hour minute sec =
  Time.UTCTime (Time.fromGregorian year month day) $
    Time.timeOfDayToTime $ Time.TimeOfDay hour minute sec

-- | Sleep until the specified time.
sleepUntil :: Time.UTCTime -> IO ()
sleepUntil !end = do
  now <- Time.getCurrentTime
  when (now < end) $ do
    let toSleep = 0.9 * realToFrac (Time.diffUTCTime end now) :: Double
    when (toSleep > 0.1) $ threadDelay $ floor $ toSleep * 1000000
    sleepUntil end

-- | Solve the first problem of the given contest and submit the solution.
solveAndSubmit :: T.Text -> IO ()
solveAndSubmit contest = do
  curl <- newCurl
  descr <- getProblemDescription curl contest
  let prog = solve descr
  T.writeFile "submitted.hs" prog
  submitProgram curl contest Haskell prog

-- | Given a problem description, create a Haskell source file that solves it.
solve :: T.Text -> T.Text
solve desc =
  makeProblem (extractInstances desc) (extractIntLiterals desc) $ \problem ->
    prettyProgram problem $ findSolution problem

----------------------------------------------------------------
-- I/O with AtCoder

-- | Create a new 'Curl' with default options set.
newCurl :: IO Curl.Curl
newCurl = do
  curl <- Curl.initialize
  Curl.setopts curl
    [ Curl.CurlCookieFile "cookies"
    , Curl.CurlCookieJar "cookies"
    , Curl.CurlUserAgent "AlmostBruteforceCoder/0.2"
    , Curl.CurlVerbose True
    ]
  return curl

-- | Get the HTML page describing the first problem of the given contest.
getProblemDescription :: Curl.Curl -> T.Text -> IO T.Text
getProblemDescription curl contest = do
  (200, content) <- getURL curl $ firstProblemPage contest
  return content

data Language = Perl | Haskell
  deriving (Show)

-- | Submit the given text as a solution to the first problem of the given
-- contest.
submitProgram :: Curl.Curl -> T.Text -> Language -> T.Text -> IO ()
submitProgram curl contest lang code = do
  (200, submitPageContents) <- getURL curl $ submitPage contest
  let !sinfo = parseSubmitInfo submitPageContents
  _ <- submitWithInfo curl contest sinfo lang code
  return ()

type URL = String
type Response = Curl.CurlResponse_ [(String, String)] S.ByteString

submitPage :: T.Text -> URL
submitPage contestName = [qc|http://{contestName}.contest.atcoder.jp/submit|]

submitPost :: T.Text -> T.Text -> URL
submitPost contestName firstPid =
  [qc|http://{contestName}.contest.atcoder.jp/submit?task_id={firstPid}|]

firstProblemPage :: T.Text -> URL
firstProblemPage contestName =
  [qc|http://{contestName}.contest.atcoder.jp/tasks/{contestName}_1|]

_loginPage :: URL
_loginPage = "http://abc010.contest.atcoder.jp/login?next_url=http%3A%2F%2Fabc010.contest.atcoder.jp%2Fsubmit#40160"

loginPost :: URL
loginPost = "http://abc010.contest.atcoder.jp/login?next_url=http%3A%2F%2Fabc010.contest.atcoder.jp%2Fsubmit"

logIn :: Curl.Curl -> IO (Int, T.Text)
logIn curl = do
  [atcoderId, atcoderPassword] <- T.lines <$> readUTF8File "atcoder_passwd"
  postURL curl loginPost
    [ ("name", atcoderId)
    , ("password", atcoderPassword)
    ]

-- | Read a UTF-8 encoded file.
readUTF8File :: FilePath -> IO T.Text
readUTF8File path = T.decodeUtf8 <$> S.readFile path

submitWithInfo
  :: Curl.Curl -> T.Text -> SubmitInfo -> Language -> T.Text
  -> IO (Int, T.Text)
submitWithInfo curl contest sinfo lang code =
    postURL curl (submitPost contest taskId) $
  [ ("__session", sessionId sinfo)
  , ("task_id", taskId)
  , ("source_code", code)
  ] ++ do
    tid <- problemIds sinfo
    return ("language_id_" <> tid, langId)
  where
    taskId = head' "no task" $ problemIds sinfo 
    langId = case lang of
      Perl -> "8"
      Haskell -> "11"

-- | Send a GET request to the given url, and return the status code and the
-- content of the page.
getURL :: Curl.Curl -> URL -> IO (Int, T.Text)
getURL curl url =
  checkResp =<< Curl.do_curl_ curl url [Curl.CurlHttpGet True]

-- | Send a POST request to the given url with the given set of parameters.
postURL :: Curl.Curl -> URL -> [(T.Text, T.Text)] -> IO (Int, T.Text)
postURL curl url params =
  checkResp =<< Curl.do_curl_ curl url [Curl.CurlPostFields pairs]
  where
    pairs = do
      (key, value) <- params
      return [qc|{enc key}={enc value}|]

    enc = HTTPTypes.urlEncode True . T.encodeUtf8

checkResp :: Response -> IO (Int, T.Text)
checkResp resp
  | Curl.respCurlCode resp /= Curl.CurlOK =
      fail [qc|bad curl status: {Curl.respCurlCode resp}|]
  | otherwise = return (Curl.respStatus resp, T.decodeUtf8 $ Curl.respBody resp)

-- | Information needed to make a POST request for submission.
data SubmitInfo = SubmitInfo
  { problemIds :: [T.Text]
  , sessionId :: T.Text
  }
  deriving (Show)

-- | Parse the submission page to extract a @SubmitInfo@.
parseSubmitInfo :: T.Text -> SubmitInfo
parseSubmitInfo txt@(TS.parseTags -> tags) = SubmitInfo
  { problemIds = pids
  , sessionId = sid
  }
  where
    sid = head' [qc|no session id found: {txt}|] $ do
      TS.TagOpen "input" attrs <- tags
      guard $ elem ("name", "__session") attrs
      ("value", value) <- attrs
      return value
    pids = head' [qc|no task selector found: {txt}|] $ do
      sec <-
        takeWhile (not . TS.tagCloseLit "select") <$>
        TS.sections (TS.tagOpenAttrLit "select" ("name", "task_id")) tags
      return $ do
        TS.TagOpen "option" attrs <- sec
        ("value", value) <- attrs
        return value

-- | 'head' with a better error message.
head' :: String -> [a] -> a
head' msg [] = error [qc|head': {msg}|]
head' _ (x:_) = x

----------------------------------------------------------------
-- Pretty-printing a program

-- | Create Haskell source text from a 'Program'.
prettyProgram :: Problem oty -> Program oty -> T.Text
prettyProgram prob prog = TL.toStrict $ TB.toLazyText builder
  where
    builder
      = TB.fromText prolog
      <> pprInputStatement (inputPattern $ problemSpec prob)
      <> "\n  "
      <> pprOutputStatement prog
      <> TB.fromText epilog

prolog :: T.Text
prolog = [q|
{-# LANGUAGE CPP #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}

import qualified Data.ByteString.Char8 as S
import Data.Monoid

main :: IO ()
main = do
  |]

epilog :: T.Text
epilog = [qc|

-- This program was generated and submitted by the following program:

# if 0

{thisProgramText}
# endif
|]

-- | The source text of this program.
thisProgramText :: String
thisProgramText = $(litE . stringL =<< runIO (readFile "abc.hs"))

-- | Format an input pattern into a statement that binds input variables.
pprInputStatement :: InputPattern -> TB.Builder
pprInputStatement ip = pat <> " <- S.getContents"
  where
    pat = evalState (pprInputPatten ip) $ PatPprState 0 0

-- | Monad for pretty-printing an input pattern.
type PatPpr = State PatPprState

data PatPprState = PatPprState
  { _usedInts :: !Int -- ^ How many int variables have we bound?
  , _usedStrs :: !Int -- ^ Ditto for string variables
  }

usedInts :: Lens' PatPprState Int
usedInts = lens _usedInts $ \s x -> s{ _usedInts = x }

usedStrs :: Lens' PatPprState Int
usedStrs = lens _usedStrs $ \s x -> s{ _usedStrs = x }

pprInputPatten :: InputPattern -> PatPpr TB.Builder
pprInputPatten GetContents = strVarPattern
pprInputPatten (Lines ls) = do
  lps <- mapM pprLinePattern ls
  return $ "(S.lines -> " <> pprList lps <> ")"

pprLinePattern :: LinePattern -> PatPpr TB.Builder
pprLinePattern GetLine = strVarPattern
pprLinePattern (Words ws) = do
  wps <- mapM pprWordPattern ws
  return $ "S.words -> " <> pprList wps

pprWordPattern :: WordPattern -> PatPpr TB.Builder
pprWordPattern Id = strVarPattern
pprWordPattern Read = do
  varPat <- intVarPattern
  return $ "S.unpack -> read -> (" <> varPat <> " :: Int)"

pprList :: [TB.Builder] -> TB.Builder
pprList list = "[" <> go list
  where
    go [] = "]"
    go (x:xs) = x <> go1 xs

    go1 [] = "]"
    go1 (x:xs) = ", " <> x <> go1 xs

strVarPattern :: PatPpr TB.Builder
strVarPattern = do
  k <- use usedStrs
  usedStrs .= k + 1
  return $ pprStrArgument k

intVarPattern :: PatPpr TB.Builder
intVarPattern = do
  k <- use usedInts
  usedInts .= k + 1
  return $ pprIntArgument k

-- | Format a program into a statement that prints the output.
pprOutputStatement :: Program oty -> TB.Builder
pprOutputStatement (IntProgram ip)
  = (const "print" `apply` pprInt ip) (-100)
pprOutputStatement (BoolProgram ip)
  = "S.putStrLn $ if " <> pprBool ip (-100) <> " then \"YES\" else \"NO\""
pprOutputStatement (StrProgram ip)
  = (const "S.putStrLn" `apply` pprStr ip) (-100)

pprInt :: IntProgram -> Int -> TB.Builder
pprInt prog = case prog of
  a :+ b -> parens 6 $ pprInt a 6 <> " + " <> pprInt b 7
  a :- b -> parens 6 $ pprInt a 6 <> " - " <> pprInt b 7
  a :* b -> parens 7 $ pprInt a 7 <> " * " <> pprInt b 8
  a :/ b -> parens 7 $ pprInt a 7 <> " `div` " <> pprInt b 8
  a :/^ b -> const "(\\x y -> (x + y - 1) `div` y)"
    `apply` pprInt a `apply` pprInt b
  a :% b -> parens 7 $ pprInt a 7 <> " `mod` " <> pprInt b 8
  Max a b -> const "max" `apply` pprInt a `apply` pprInt b
  Min a b -> const "min" `apply` pprInt a `apply` pprInt b
  NumChars s -> const "S.length" `apply` pprStr s
  NumLines s -> const "((+1) . S.count '\n')" `apply` pprStr s
  LiteralInt n -> const $ TB.decimal n
  IntArgument k -> const $ pprIntArgument k

pprStr :: StrProgram -> Int -> TB.Builder
pprStr prog = case prog of
  a :++ b -> parens 6 $ pprStr a 7 <> " <> " <> pprStr b 6
  Show a -> const "S.pack" `apply` (const "show" `apply` pprInt a)
  Sep a b -> parens 6 $ pprStr a 7 <> " <> \" \" <> "  <> pprStr b 6
  VCat a b -> parens 6 $ pprStr a 7 <> " <> \"\\n\" <> "  <> pprStr b 6
  LiteralString str -> const $ TB.fromString (show str)
  StrArgument k -> const $ pprStrArgument k

pprBool :: BoolProgram -> Int -> TB.Builder
pprBool prog = case prog of
  a :&& b -> parens 3 $ pprBool a 4 <> " && " <> pprBool b 3
  a :|| b -> parens 2 $ pprBool a 3 <> " && " <> pprBool b 2
  Not a -> const "not" `apply` pprBool a
  a :< b -> parens 4 $ pprInt a 5 <> " < " <> pprInt b 5
  a :== b -> parens 4 $ pprInt a 5 <> " == " <> pprInt b 5

pprIntArgument :: Int -> TB.Builder
pprIntArgument k = "i" <> TB.decimal k

pprStrArgument :: Int -> TB.Builder
pprStrArgument k = "s" <> TB.decimal k

apply :: (Int -> TB.Builder) -> (Int -> TB.Builder) -> Int -> TB.Builder
apply f x = parens 10 $ f 10 <> " " <> x 11

parens :: Int -> TB.Builder -> Int -> TB.Builder
parens k b prec
  | prec > k = "(" <> b <> ")"
  | otherwise = b

----------------------------------------------------------------
-- Finding a solution

-- | Return a program that passes all the test cases. If there is no such
-- representable program, it loops forever.
findSolution :: (ObjectType oty) => Problem oty -> Program oty
findSolution prob = head $ filter good $ generateAll prob
  where
    good prog = all (satisfiesExample prog) $ instances prob

-- | Returns an inifite list that enumerates (almost) all representable
-- programs.
generateAll :: (ObjectType oty) => Problem oty -> [Program oty]
generateAll prob = concatMap gen [0..]
  where
    spec = problemSpec prob
    ints = intsInSpec prob
    gen sz = generateProgramsWithSize spec ints sz
      ++ do
        guard $ sz' >= 0
        (prob', embed) <- reds
        map embed $ generateProgramsWithSize (problemSpec prob') ints sz'
      where
        sz' = sz - 2
    reds = reductions prob

-- | Test whether the given program behaves in the way that is specified
-- in an example.
satisfiesExample :: (ObjectType oty) => Program oty -> TypedInstance oty -> Bool
satisfiesExample prog (TypedInstance inputs expected) =
  eval prog inputs == Just expected

-- | Evaluate the given program with the given inputs. Returns the outcmome
-- of the evaluation, or @Nothing@ if an error occurs.
eval :: Program oty -> TypedInput -> Maybe oty
eval prog input = runReaderT go input
  where
    go = case prog of
      IntProgram p -> evalInt p
      StrProgram p -> evalString p
      BoolProgram p -> evalBool p

-- | A monad for evaluating a program.
type Eval = ReaderT TypedInput Maybe

evalString :: StrProgram -> Eval S.ByteString
evalString prog = seqM $ case prog of
  a :++ b -> (<>) <$> evalString a <*> evalString b
  Sep a b -> do
    x <- evalString a
    y <- evalString b
    return $! x <> " " <> y
  VCat a b -> do
    x <- evalString a
    y <- evalString b
    return $! x <> "\n" <> y
  Show a -> S.pack . show <$> evalInt a
  LiteralString a -> return a
  StrArgument k -> asks $ (V.!k) . strInput

evalInt :: IntProgram -> Eval Int
evalInt prog = seqM $ case prog of
  a :+ b -> (+) <$> evalInt a <*> evalInt b
  a :- b -> (-) <$> evalInt a <*> evalInt b
  a :* b -> (*) <$> evalInt a <*> evalInt b
  a :/ b -> div <$> evalInt a <*> (nonzero =<< evalInt b)
  a :/^ b -> cdiv <$> evalInt a <*> (nonzero =<< evalInt b)
  a :% b -> mod <$> evalInt a <*> (nonzero =<< evalInt b)
  Max a b -> max <$> evalInt a <*> evalInt b
  Min a b -> min <$> evalInt a <*> evalInt b
  NumChars a -> S.length <$> evalString a
  NumLines a -> (+1) . S.count '\n' <$> evalString a
  LiteralInt n -> return n
  IntArgument k -> asks $ (U.!k) . intInput
  where
    cdiv a b = (a+b-1) `div` b

    nonzero 0 = mzero
    nonzero x = return x

evalBool :: BoolProgram -> Eval Bool
evalBool prog = seqM $ case prog of
  a :&& b -> ifM (evalBool a) (evalBool b) (return False)
  a :|| b -> ifM (evalBool a) (return True) (evalBool b)
  Not a -> not <$> evalBool a
  a :< b -> (<) <$> evalInt a <*> evalInt b
  a :== b -> (==) <$> evalInt a <*> evalInt b

seqM :: (Monad m) => m a -> m a
seqM = (>>= (return$!))
{-# INLINE seqM #-}

ifM :: (Monad m) => m Bool -> m a -> m a -> m a
ifM a t f = a >>= \case
  True -> t
  False -> f
{-# INLINE ifM #-}

----------------------------------------------------------------
-- Reducing a problem

-- | Try to reduce the problem a bit by looking for a common pattern
-- in example outputs. Returns a list of possible reductions. Each reduction
-- is a pair: the first element is the reduced
-- problem and the second element is a function that takes a solution
-- to the reduced problem and produces a solution to the original problem.
--
-- For example, if all the sample outputs end with a "pp" suffix, it generates
-- a "reduced" problem whose task is to generate the output string without the
-- "pp" suffix.
reductions
  :: forall oty
  . (ObjectType oty)
  => Problem oty
  -> [(Problem oty, Program oty -> Program oty)]
reductions prob = do
  modifier <- modifiers
  (outputs, embed) <- modifier $ map expectedOutput iss
  let newProb = prob{ instances = zipWith replaceOutput iss outputs }
  return (newProb, embed)
  where
    iss = instances prob
    replaceOutput inst out = inst{ expectedOutput = out }

    modifiers :: [[oty] -> [([oty], Program oty -> Program oty)]]
    modifiers = case objectType :: ObjectTypeTag oty of
      Bool -> []
      String -> [prefix, suffix]
      Int -> [add, mul]

    prefix = strip commonPrefixes S.drop (:++)
    suffix = strip commonSuffixes (\n s -> S.take (S.length s-n) s) (flip (:++))

    add, mul :: [Int] -> [([Int], Program Int -> Program Int)]
    add xs = do
      guard $ m /= 0
      let
        embed (IntProgram s) = IntProgram (s :+ LiteralInt m)
      return (map (subtract m) xs, embed)
      where
        m = minimum xs

    mul xs = do
      guard $ g /= 0 && g /= 1
      let
        embed (IntProgram s) = IntProgram (s :* LiteralInt g)
      return (map (flip div g) xs, embed)
      where
        g = foldl1' gcd xs

    strip
      :: ([S.ByteString] -> [S.ByteString])
      -> (Int -> S.ByteString -> S.ByteString)
      -> (StrProgram -> StrProgram -> StrProgram)
      -> [S.ByteString]
      -> [([S.ByteString], Program S.ByteString -> Program S.ByteString)]
    strip common f g (ss :: [S.ByteString]) = do
      c <- common ss
      guard $ not $ S.null c
      let
        embed (StrProgram s) = StrProgram (g (LiteralString c) s)
      return (map (f $ S.length c) ss, embed)

-- | Returns the list of all strings that are a prefix of all given strings.
--
-- When the empty list is given, returns @[""]@.
commonPrefixes :: [S.ByteString] -> [S.ByteString]
commonPrefixes = S.inits . commonPrefix

commonPrefix :: [S.ByteString] -> S.ByteString
commonPrefix [] = S.empty
commonPrefix xs = foldl1' common xs
  where
    common x y = S.take len x
      where
        len = length $ takeWhile id $ S.zipWith (==) x y

commonSuffixes :: [S.ByteString] -> [S.ByteString]
commonSuffixes = map S.reverse . commonPrefixes . map S.reverse

----------------------------------------------------------------
-- Generating programs

-- | A program that produces a value of type @oty@.
data Program oty where
  IntProgram :: IntProgram -> Program Int
  StrProgram :: StrProgram -> Program S.ByteString
  BoolProgram :: BoolProgram -> Program Bool

deriving instance Show (Program oty)

-- | A program that returns an Int.
data IntProgram
  = !IntProgram :+ !IntProgram
  | !IntProgram :- !IntProgram
  | !IntProgram :* !IntProgram
  | !IntProgram :/ !IntProgram -- ^ floor after division
  | !IntProgram :/^ !IntProgram -- ^ ceiling after division
  | !IntProgram :% !IntProgram
  | Max !IntProgram !IntProgram
  | Min !IntProgram !IntProgram
  | NumChars !StrProgram
  | NumLines !StrProgram
  | LiteralInt !Int
  | IntArgument !Int
  deriving (Show)

-- | A program that returns a Bool.
data BoolProgram
  = !BoolProgram :&& !BoolProgram
  | !BoolProgram :|| !BoolProgram
  | Not !BoolProgram
  | !IntProgram :< !IntProgram
  | !IntProgram :== !IntProgram
  deriving (Show)

-- | A program that returns a String.
data StrProgram
  = !StrProgram :++ !StrProgram -- ^ concatenation
  | Sep !StrProgram !StrProgram -- ^ concatenation with a space
  | VCat !StrProgram !StrProgram -- ^ concatenation with a newline
  | Show !IntProgram
  | LiteralString !S.ByteString
  | StrArgument !Int
  deriving (Show)

-- | A program-generation monad.
type Generate = ReaderT GenEnv (StateT GenState [])

-- | An argument to a program-generation process.
data GenEnv = GenEnv
  { _targetProblem :: !ProblemSpec
  , _size :: !Int
  , _allowedIntLiterals :: ![Int]
  }

targetProblem :: Lens' GenEnv ProblemSpec
targetProblem = lens _targetProblem $ \s x -> s{ _targetProblem = x }

size :: Lens' GenEnv Int
size = lens _size $ \s x -> s{ _size = x }

allowedIntLiterals :: Lens' GenEnv [Int]
allowedIntLiterals = lens _allowedIntLiterals $ \s x -> s{ _allowedIntLiterals = x }

-- | A state of a program-generation process.
data GenState = GenState
  { _intVars :: !IS.IntSet -- ^ set of used int inputs
  , _strVars :: !IS.IntSet -- ^ set of used string inputs
  }

intVars :: Lens' GenState IS.IntSet
intVars = lens _intVars $ \s x -> s{ _intVars = x }

strVars :: Lens' GenState IS.IntSet
strVars = lens _strVars $ \s x -> s{ _strVars = x }

generateProgramsWithSize
  :: forall oty. (ObjectType oty)
  => ProblemSpec -> IS.IntSet -> Int -> [Program oty]
generateProgramsWithSize spec ints sz = case objectType :: ObjectTypeTag oty of
  Int -> map IntProgram $ gen1 intProgram
  String -> map StrProgram $ gen1 strProgram
  Bool -> map BoolProgram $ gen1 boolProgram
  where
    gen1 :: Generate a -> [a]
    gen1 g = runGenerate spec ints sz g

-- | Run the given program generator to generate programs of the given size.
runGenerate :: ProblemSpec -> IS.IntSet -> Int -> Generate a -> [a]
runGenerate prob ints sz
  = map fst
  . filter (full . snd)
  . flip runStateT initialState
  . flip runReaderT env
  where
    env = GenEnv prob sz lits
    initialState = GenState IS.empty IS.empty
    lits = IS.toList $ IS.delete 0 $ IS.fromList $ do
      int <- IS.toList ints
      [int, int + 1]

    -- Filter out programs that don't use all the input variables.
    full st
      = st^.intVars.to IS.size == nIntInputs prob
      && st^.strVars.to IS.size == nStringInputs prob

-- | Generate @IntProgram@s.
intProgram :: Generate IntProgram
intProgram = msum
  [ binIntOpSym (:+)
  , binIntOp (:-)
  , binIntOpSym (:*)
  , binIntOp (:/)
  , binIntOp (:/^)
  , binIntOp (:%)
  , binIntOpSym Min
  , binIntOpSym Max
  , unOp NumChars strProgram
  , unOp NumLines strProgram
  , literal
  , argument nIntInputs intVars IntArgument
  ]
  where
    binIntOp op = binOp op intProgram intProgram
    binIntOpSym op = binOpSym op intProgram intProgram
    literal = do
      1 <- view size
      ints <- view allowedIntLiterals
      msum $ map (return . LiteralInt) ints

boolProgram :: Generate BoolProgram
boolProgram = msum
  [ binOpSym (:&&) boolProgram boolProgram
  , binOpSym (:||) boolProgram boolProgram
  , unOp Not boolProgram
  , binOp (:<) intProgram intProgram
  , binOpSym (:==) intProgram intProgram
  ]

strProgram :: Generate StrProgram
strProgram = msum
  [ binOp (:++) strProgram strProgram
  , binOp Sep strProgram strProgram
  , binOp VCat strProgram strProgram
  , unOp Show intProgram
  , argument nStringInputs strVars StrArgument
  ]

argument
  :: (ProblemSpec -> Int)
  -> Lens' GenState IS.IntSet
  -> (Int -> a)
  -> Generate a
argument inputCount vars f = do
  sz <- view size
  guard $ sz <= 2
  nvars <- view $ targetProblem.to inputCount
  varId <- msum $ map return [0..nvars-1]
  used <- IS.member varId <$> use vars

  -- Impose a bigger cost when using an already-used variable.
  guard $ sz == if used then 2 else 1
  vars %= IS.insert varId
  return $ f varId

unOp :: (a -> b) -> Generate a -> Generate b
unOp op genChild = do
  sz <- view size
  guard $ 2 <= sz
  local (size .~ sz - 1) $ op <$> genChild

binOpSym :: (a -> b -> c) -> Generate a -> Generate b -> Generate c
binOpSym = binOpMaybeSym True

binOp :: (a -> b -> c) -> Generate a -> Generate b -> Generate c
binOp = binOpMaybeSym False

binOpMaybeSym :: Bool -> (a -> b -> c) -> Generate a -> Generate b -> Generate c
binOpMaybeSym symmetric op leftChild rightChild = do
  !sz <- view size
  guard $ 3 <= sz
  let !ub = if symmetric then div sz 2 else sz - 2
  !leftSize <- msum $ map return [1 .. ub]
  let !rightSize = sz - leftSize - 1
  op
    <$> local (size .~ leftSize) leftChild
    <*> local (size .~ rightSize) rightChild

----------------------------------------------------------------
-- Guessing Input/Output formats of a problem

-- | A problem with a defined I/O format.
data Problem otype = Problem
  { problemSpec :: !ProblemSpec
  , instances :: ![TypedInstance otype]
  , intsInSpec :: !IS.IntSet -- ^ integer literals found in the spec
  }
  deriving (Show)

-- | The I/O specification of a problem, i.e. the number of arguments,
-- the type of each argument, etc.
data ProblemSpec = ProblemSpec
  { nIntInputs :: !Int
  , nStringInputs :: !Int
  , inputPattern :: !InputPattern
  , outputPattern :: !OutputPattern
  }

deriving instance Show ProblemSpec

-- | A problem instance in the parsed form.
data TypedInstance otype = TypedInstance
  { typedInput :: !TypedInput
  , expectedOutput :: !otype
  }
  deriving (Show)

-- | A concrete input to a problem.
data TypedInput = TypedInput
  { intInput :: !(U.Vector Int)
  , strInput :: !(V.Vector S.ByteString)
  }
  deriving (Show)

-- | A type of values that can be dealt with our program. Currently
-- it's one of Int, Bool or String.
class (Show a, Eq a) => ObjectType a where
  objectType :: ObjectTypeTag a

data ObjectTypeTag a where
  Int :: ObjectTypeTag Int
  String :: ObjectTypeTag S.ByteString
  Bool :: ObjectTypeTag Bool

deriving instance Show (ObjectTypeTag a)
deriving instance Eq (ObjectTypeTag a)

instance ObjectType Int where
  objectType = Int

instance ObjectType S.ByteString where
  objectType = String

instance ObjectType Bool where
  objectType = Bool

-- | How an input is encoded as a string, or an instruction to
-- systematically parse the input string into values.
--
-- Some patterns are more structured than others. For example,
-- any pattern of the form @(Lines x)@ is more structured than
-- @GetContents@.
data InputPattern
  = Lines [LinePattern] -- ^ Parse each line individually
  | GetContents -- ^ Take the whole input as a string
  deriving (Show)

-- | Instruction to parse a line.
data LinePattern
  = Words [WordPattern] -- ^ Parse each space-separated word individually
  | GetLine -- ^ Take the whole line as a string
  deriving (Show)

-- | Instruction to parse a space-separated word.
data WordPattern
  = Read -- ^ Read the word as an integer
  | Id -- ^ Take the word as a string
  deriving (Show)

-- | Instruction to parse an output example. Let's assume that the output
-- always consists of a single value
data OutputPattern = forall a. (ObjectType a) =>
  OutputPattern (ObjectTypeTag a)

deriving instance Show OutputPattern
instance Eq OutputPattern where
  OutputPattern Int == OutputPattern Int = True
  OutputPattern Bool == OutputPattern Bool = True
  OutputPattern String == OutputPattern String = True
  _ == _ = False

-- | Create a problem from examples.
makeProblem
  :: [Instance]
  -> IS.IntSet
  -> (forall oty. (ObjectType oty) => Problem oty -> r)
  -> r
makeProblem is ints cont =
  parseInstances (inputPattern spec) (outputPattern spec) is $ \ts ->
    cont $ Problem spec ts ints
  where
    spec = guessProblemSpec is

-- | Guess the I/O specificaitons using examples.
guessProblemSpec :: [Instance] -> ProblemSpec
guessProblemSpec is = ProblemSpec{..}
  where
    (nIntInputs, nStringInputs) = countInputs inputPattern
    inputPattern = lubs "guessInputPattern: no instance" $
      map (bestInputPattern . fst) is
    outputPattern = lubs "guessInputPattern: no instance" $
      map (bestOutputPattern . snd) is

-- | Parse the given instances using the given I/O patterns. This function is
-- in the CPS in order to avoid an explicit use of existential types.
parseInstances
  :: InputPattern
  -> OutputPattern
  -> [Instance]
  -> (forall otype. (ObjectType otype) => [TypedInstance otype] -> r)
  -> r
parseInstances ipat (OutputPattern (_ :: ObjectTypeTag oty)) is cont
  = cont (map (parseInstance ipat) is :: [TypedInstance oty])

-- | Parse one problem instance according to the input pattern.
-- The output pattern is inferred from the type.
parseInstance
  :: (ObjectType otype)
  => InputPattern -> Instance -> TypedInstance otype
parseInstance ipat (input, output) = TypedInstance
  { typedInput = TypedInput
    { intInput = U.fromList ints
    , strInput = V.fromList strs
    }
  , expectedOutput = parseOutput output
  }
  where
    (ints, strs) = parseInput ipat input

-- | Parse an input according to the input pattern.
parseInput :: InputPattern -> T.Text -> ([Int], [S.ByteString])
parseInput (Lines ps) txt = mconcat $ zipWith parseLine ps $ T.lines txt
parseInput GetContents txt = ([], [readString txt])

parseLine :: LinePattern -> T.Text -> ([Int], [S.ByteString])
parseLine (Words ps) txt = mconcat $ zipWith parseWord ps $ T.words txt
parseLine GetLine txt = ([], [readString txt])

parseWord :: WordPattern -> T.Text -> ([Int], [S.ByteString])
parseWord Read t = ([readInt t], [])
parseWord Id t = ([], [readString t])

parseOutput :: forall otype. (ObjectType otype) => T.Text -> otype
parseOutput txt = case objectType :: ObjectTypeTag otype of
  Int -> readInt txt
  String -> readString txt
  Bool -> fromMaybe (error [qc|cannot read bool: {show txt}|]) $
    readBoolMaybe txt

readString :: T.Text -> S.ByteString
readString = T.encodeUtf8

readInt :: T.Text -> Int
readInt txt = fromMaybe (error [qc|cannot read int: {show txt}|]) $
  readIntMaybe txt

readIntMaybe :: T.Text -> Maybe Int
readIntMaybe txt = case T.signed T.decimal txt of
  Right (val, rest)
    | T.null rest -> Just val
  _ -> Nothing

readBoolMaybe :: T.Text -> Maybe Bool
readBoolMaybe "YES" = Just True
readBoolMaybe "NO" = Just False
readBoolMaybe _ = Nothing

-- | Count the number of integer and string inputs in a pattern.
countInputs :: InputPattern -> (Int, Int)
countInputs (Lines ls) = sumV $ map countInputsInLine ls
countInputs GetContents = (0, 1)

countInputsInLine :: LinePattern -> (Int, Int)
countInputsInLine (Words ws) = sumV $ map countInputsInWord ws
countInputsInLine GetLine = (0, 1)

countInputsInWord :: WordPattern -> (Int, Int)
countInputsInWord Read = (1, 0)
countInputsInWord Id = (0, 1)

-- | Semilattice, or commutative idempotent semigroup.
class Semilattice a where
  lub :: a -> a -> a

-- The {Input,Line,Word}Pattern types each forms a semilattice where
-- @(a `lub` b)@ means the "most structured" pattern that is compatible
-- with both @a@ and @b@.
instance Semilattice InputPattern where
  lub (Lines xs) (Lines ys) =
    maybe GetContents Lines $ lubList xs ys
  lub _ _ = GetContents

instance Semilattice LinePattern where
  lub (Words xs) (Words ys) = 
    maybe GetLine Words $ lubList xs ys
  lub _ _ = GetLine

instance Semilattice WordPattern where
  lub Read Read = Read
  lub _ _ = Id

instance Semilattice OutputPattern where
  lub a b
    | a == b = a
  lub _ _ = OutputPattern String

-- | Take the lub of the given values.
lubs :: (Semilattice a) => String -> [a] -> a
lubs msg [] = error msg
lubs _ is = foldl1' lub is

-- | Pairwise lub, if the lists have the same length.
lubList :: (Semilattice a) => [a] -> [a] -> Maybe [a]
lubList xs ys
  | length xs == length ys = Just $ zipWith lub xs ys
  | otherwise = Nothing

-- | Return the most structured input pattern that can parse the
-- given text.
bestInputPattern :: T.Text -> InputPattern
bestInputPattern = Lines . map bestLinePattern . T.lines

bestLinePattern :: T.Text -> LinePattern
bestLinePattern = Words . map bestWordPattern . T.words

bestWordPattern :: T.Text -> WordPattern
bestWordPattern txt
  | isJust (readIntMaybe txt) = Read
  | otherwise = Id

bestOutputPattern :: T.Text -> OutputPattern
bestOutputPattern txt
  | isJust (readIntMaybe txt) = OutputPattern Int
  | isJust (readBoolMaybe txt) = OutputPattern Bool
  | otherwise = OutputPattern String

----------------------------------------------------------------
-- Extracting problem instances from HTML

-- | A pair of sample input and output.
type Instance = (T.Text, T.Text)

-- | Extract problem instances from a problem page.
extractInstances :: T.Text -> [Instance]
extractInstances
  = combineSamples
  . map extractSample
  . TS.partitions (isJust . parseSampleHead)
  . TS.parseTags

-- | A sample input or output.
data Sample = Sample
  { _inputOutput :: !InputOutput
  , _instanceName :: !T.Text
  , _sampleBody :: !T.Text
  }

data InputOutput = Input | Output
  deriving (Eq)

-- | Try to parse the given tag as an indicator of a sample.
parseSampleHead :: TS.Tag T.Text -> Maybe (InputOutput, T.Text)
parseSampleHead (TS.TagText (T.strip -> txt)) =
  ((Input,) <$> T.stripPrefix "入力例" txt) <|>
  ((Output,) <$> T.stripPrefix "出力例" txt)
parseSampleHead _ = Nothing

-- | Extract one sample from HTML.
extractSample :: [TS.Tag T.Text] -> Sample
extractSample
  ((parseSampleHead -> Just (io, name)) : ts0)
    = case TS.partitions (TS.tagOpenNameLit "pre") ts0 of
      [] -> error "extractSample: no pre found"
      ts1 : _ -> fromMaybe (error "extractSample: no text found") $
        fmap (Sample io name . T.strip) $ msum $ map TS.maybeTagText ts1
extractSample ts = error [qc|extractSample: no head: {take 1 ts}|]

-- | Combine matching input and output samples to build instances.
combineSamples :: [Sample] -> [Instance]
combineSamples samples
  = mapMaybe buildInstance
  $ M.toList
  $ M.fromListWith (++)
  [ (name, [(io, body)]) | Sample io name body <- samples]
  where
    buildInstance (name, ss) = case partition ((==Input) . fst) ss of
      ([(_, input)], [(_, output)]) -> Just (input, output)
      ([], _) -> warn "no input"
      (_, []) -> warn "no output"
      (_:_:_, _) -> warn "multiple inputs"
      (_, _:_:_) -> warn "multiple outputs"
      where
        warn msg = trace [qc|combineSamples: {name}: {msg::String}|] Nothing

----------------------------------------------------------------
-- Extracting integer literals from HTML

-- | Extract integer literals from a problem page.
extractIntLiterals :: T.Text -> IS.IntSet
extractIntLiterals
  = IS.fromList
  . concatMap extractIntsFromText
  . mapMaybe TS.maybeTagText
  . interestingRegion
  . TS.parseTags
  where
    interestingRegion
      = snd . break (TS.tagOpenAttrLit "div" ("id", "task-statement"))
      . fst . break (TS.tagText $ (=="入力例1") . T.strip)

extractIntsFromText :: T.Text -> [Int]
extractIntsFromText text = do
  grp <- T.groupBy ((==) `on` isAsciiDigit) text
  maybeToList $ readIntMaybe grp
  where
    isAsciiDigit x = '0' <= x && x <= '9'

# endif

Submission

Task問題 A - 閉路グラフ
User nameユーザ名 mkotha
Created time投稿日時
Language言語 Haskell (GHC 7.4.1)
Status状態 AC
Score得点 100
Source lengthソースコード長 34770 Byte
File nameファイル名
Exec time実行時間 31 ms
Memory usageメモリ使用量 1432 KB

Test case

Set

Set name Score得点 / Max score Cases
Sample - subtask0_sample_01.txt,subtask0_sample_02.txt,subtask0_sample_03.txt,subtask0_sample_04.txt
All 100 / 100 subtask0_sample_01.txt,subtask0_sample_02.txt,subtask0_sample_03.txt,subtask0_sample_04.txt,subtask1_01.txt,subtask1_02.txt,subtask1_03.txt,subtask1_04.txt,subtask1_05.txt,subtask1_06.txt,subtask1_07.txt,subtask1_08.txt,subtask1_09.txt,subtask1_10.txt,subtask1_11.txt,subtask1_12.txt,subtask1_13.txt,subtask1_14.txt,subtask1_15.txt,subtask1_16.txt,subtask1_17.txt,subtask1_18.txt,subtask1_19.txt,subtask1_20.txt,subtask1_21.txt,subtask1_22.txt,subtask1_23.txt,subtask1_24.txt,subtask1_25.txt,subtask1_26.txt

Test case

Case name Status状態 Exec time実行時間 Memory usageメモリ使用量
subtask0_sample_01.txt AC 28 ms 1312 KB
subtask0_sample_02.txt AC 31 ms 1380 KB
subtask0_sample_03.txt AC 30 ms 1316 KB
subtask0_sample_04.txt AC 29 ms 1312 KB
subtask1_01.txt AC 31 ms 1308 KB
subtask1_02.txt AC 30 ms 1364 KB
subtask1_03.txt AC 29 ms 1396 KB
subtask1_04.txt AC 30 ms 1396 KB
subtask1_05.txt AC 28 ms 1308 KB
subtask1_06.txt AC 29 ms 1384 KB
subtask1_07.txt AC 30 ms 1364 KB
subtask1_08.txt AC 30 ms 1304 KB
subtask1_09.txt AC 31 ms 1384 KB
subtask1_10.txt AC 28 ms 1392 KB
subtask1_11.txt AC 30 ms 1432 KB
subtask1_12.txt AC 30 ms 1312 KB
subtask1_13.txt AC 28 ms 1432 KB
subtask1_14.txt AC 28 ms 1308 KB
subtask1_15.txt AC 28 ms 1308 KB
subtask1_16.txt AC 30 ms 1368 KB
subtask1_17.txt AC 29 ms 1368 KB
subtask1_18.txt AC 29 ms 1396 KB
subtask1_19.txt AC 30 ms 1312 KB
subtask1_20.txt AC 30 ms 1308 KB
subtask1_21.txt AC 28 ms 1384 KB
subtask1_22.txt AC 29 ms 1308 KB
subtask1_23.txt AC 27 ms 1312 KB
subtask1_24.txt AC 29 ms 1304 KB
subtask1_25.txt AC 29 ms 1304 KB
subtask1_26.txt AC 29 ms 1368 KB