AtCoder の新環境テストの問題を解いてみる

#include "../src/Header.hs"

main :: IO ()
main = do
  [x] <- readLineInputs @Int
  [y, z] <- readLineInputs @Int
  s <- BS.getLine

  printN $ x + y + z
  putSpace
  BS.putStrLn s

#include "../src/Header.hs"

main :: IO ()
main = do
  [x, y] <- readLineInputs

  putStrLn $ solve x y

solve :: Int -> Int -> String
solve x y
  | (x * y) `mod` 2 == 0 = "Even"
  | otherwise            = "Odd"

#include "../src/Header.hs"

main :: IO ()
main = do
  s <- Text.getLine
  print $ ocount (== '1') s

#include "../src/Header.hs"

main :: IO ()
main = do
  discardLine
  xs <- readLineInputs

  print $ minimum [ checkShifts x | x <- xs ]

-- |
--
-- >>> checkShifts 0
-- 0
-- >>> checkShifts 2
-- 1
-- >>> checkShifts 20
-- 2
--
checkShifts :: Int -> Int
checkShifts = go 0
  where
    go m 0 = m
    go m n
      | n .&. 1 == 1 = m
      | otherwise    = go (m + 1) $ n `shiftR` 1

#include "../src/Header.hs"

main :: IO ()
main = do
  [a] <- readLineInputs
  [b] <- readLineInputs
  [c] <- readLineInputs
  [x] <- readLineInputs

  print $ solve a b c x

-- |
--
-- >>> solve 1 1 1 0
-- 1
--
solve :: Int -> Int -> Int -> Int -> Int
solve a b c x = length @[] do
  a' <- [0..a]
  b' <- [0..b]
  let x' = x - a' * 500 - b' * 100
  withFilter $ x' >= 0 && x' <= c * 50
  pure ()

#include "../src/Header.hs"

main :: IO ()
main = do
  [n, a, b] <- getLineInputs @Int

  print $ sum [ x | x <- [1..n], let k = sumDigits x, a <= k, k <= b ]

-- |
--
-- >>> sumDigits 11
-- 2
--
sumDigits :: Int -> Int
sumDigits = go 0
  where
    go m 0 = m
    go m n = go (m + n `mod` 10) $ n `div` 10

#include "../src/Header.hs"

main :: IO ()
main = do
  discardLine
  xs <- readLineInputsMVec @Int

  r <- solve xs
  print r

-- |
--
-- >>> Vec.unsafeThaw [3, 1] >>= solve
-- 2
-- >>> Vec.unsafeThaw [20, 18, 2, 18] >>= solve
-- 18
--
solve :: MVector Int -> IO Int
solve ys = do
    MVec.sortBy rcompare ys
    ifoldlMVec' go 0 ys
  where
    go m i x = m + if i .&. 1 == 0 then x else negate x

#include "../src/Header.hs"

main :: IO ()
main = do
  [n] <- readLineInputs @Int
  xs :: MVector Int <- MVec.replicateM n do
    [x] <- readLineInputs
    pure x

  MVec.sort xs
  xs' <- Vec.unsafeFreeze xs
  print $ olength $ Vec.uniq xs'

#include "../src/Header.hs"

main :: IO ()
main = do
  [n, y] <- readLineInputs @Int

  case solve n y of
    Nothing           -> putStrLn "-1 -1 -1"
    Just (i1, i2, i3) -> prints [i1, i2, i3]

-- |
--
-- prop> maybe n (\(x, y, z) -> x + y + z) (solve n $ m * 1000) == n
-- >>> solve 9 45000
-- Just ...
-- >>> solve 20 196
-- Nothing
--
solve :: Int -> Int -> Maybe (Int, Int, Int)
solve n y = headMay @[_] do
  i1 <- [0..n]
  let n2 = n - i1
  i2 <- [0..n2]
  let i3 = n2 - i2
  let y' = y - i1 * 10000 - i2 * 5000
  withFilter $ y' == i3 * 1000
  pure (i1, i2, i3)

#include "../src/Header.hs"

main :: IO ()
main = do
  s <- BS.getLine

  putStrLn case solve s of
    True  -> "YES"
    False -> "NO"

-- |
--
-- >>> solve "erasedream"
-- True
--
solve :: ByteString -> Bool
solve s = case parseOnlyEof p $ BS.reverse s of
    Left{}  -> False
    Right{} -> True
  where
    p = Parse.skipMany $ altconcat [ Parse.try $ Parse.string $ BS.reverse w | w <- ws ]

    ws = [ "dream", "dreamer", "erase", "eraser" ]

#include "../src/Header.hs"

main :: IO ()
main = do
  [n] <- readLineInputs
  solve n

solve :: Int -> IO ()
solve m = go m 0 0 0
  where
    go :: Int -> Int -> Int -> Int -> IO ()
    go n t x y
      | n == 0    = putStrLn "Yes"
      | otherwise = do
        [t', x', y'] <- readLineInputs
        let d = abs (x - x') + abs (y - y')
        let td = t' - t
        case td >= d && (td - d) `mod` 2 == 0 of
          True  -> go (n - 1) t' x' y'
          False -> do
            replicateM_ (n - 1) discardLine
            putStrLn "No"

#include "../src/Header.hs"

main :: IO ()
main = do
  [n, _] <- readLineInputs
  solve n $ Vec.iterateN n succ 'A'

solve :: Int -> Vector Char -> IO ()
solve n v = do
    xs1 <- Vec.unsafeThaw v
    case n of
      5  -> go1 xs1
      26 -> go2 xs1
      _  -> error $ "unexpected input:" ++ show n
    xs2 <- Vec.unsafeFreeze xs1

    putStr "! "
    oforM_ xs2 putChar
    putEndLine
  where
    go1 xs = do
      askSwap xs 0 1
      askSwap xs 2 3
      ask xs 0 2 >>= \case
        True  -> pure ()
        False -> do
          MVec.swap xs 0 2
          MVec.swap xs 1 3
      --   2 - 3
      --  /
      -- 0 - 1

      ask xs 2 4 >>= \case
        --     4
        --    /
        --   2 - 3
        --  /
        -- 0 - 1
        True  -> askSwap xs 3 4
        -- 4
        --  \
        --   2 - 3
        --  /
        -- 0 - 1
        False -> do
          askSwap xs 0 4
          MVec.swap xs 4 2
          MVec.swap xs 4 3
      -- 0 - 2 - 3 - 4
      --  \
      --   1


      ask xs 1 3 >>= \case
        -- 0 - 2 - 3 - 4
        --  \     /
        --   1 --
        True  -> askSwap xs 1 2
        -- 0 - 2 - 3 - 4
        --          \
        --           1
        False -> do
          askSwap xs 1 4
          MVec.swap xs 3 1
          MVec.swap xs 2 1

    go2 xs = sortM askChar xs

    askSwap xs i1 i2 = ask xs i1 i2 >>= \case
      True  -> pure ()
      False -> MVec.swap xs i1 i2

    ask xs i1 i2 = do
      c1 <- MVec.read xs i1
      c2 <- MVec.read xs i2
      askChar c1 c2

    askChar c1 c2 = do
      putStrLn ['?', ' ', c1, ' ', c2]
      putFlush
      [c] <- getLine
      pure case c of
        '<' -> True
        '>' -> False
        _   -> error $ "unexpected input: " ++ [c]

-- |
--
-- >>> xs <- mvecFromList ['B','A','E','F','D','C']
-- >>> sortM (\x y -> pure $ x < y) xs
-- >>> mvecToList xs
-- "ABCDEF"
--
sortM :: (Char -> Char -> IO Bool) -> MVector Char -> IO ()
sortM cmp v = do
    v2 <- MVec.clone v
    go 0 (MVec.length v) v2 v
  where
    go i l xs ys = if
      | l == 1    -> pure ()
      | otherwise -> do
        let i1 = i
            l1 = l `div` 2
        go i1 l1 ys xs
        let i2 = i + l1
            l2 = l - l1
        go i2 l2 ys xs

        merge i1 i1 l1 i2 l2 xs ys

    merge i i1 l1 i2 l2 xs ys = if
      | l1 == 0 -> do
        let xs2 = MVec.slice i2 l2 xs
        let ys2 = MVec.slice i l2 ys
        MVec.copy ys2 xs2
      | l2 == 0 -> do
        let xs1 = MVec.slice i1 l1 xs
        let ys1 = MVec.slice i l1 ys
        MVec.copy ys1 xs1
      | otherwise -> do
        c1 <- MVec.read xs i1
        c2 <- MVec.read xs i2
        cmp c1 c2 >>= \case
          True -> do
            MVec.write ys i c1
            merge (i + 1) (i1 + 1) (l1 - 1) i2 l2 xs ys
          False -> do
            MVec.write ys i c2
            merge (i + 1) i1 l1 (i2 + 1) (l2 - 1) xs ys