結果

問題 No.2356 Back Door Tour in Four Seasons
コンテスト
ユーザー mtwtkman
提出日時 2023-06-18 13:31:01
言語 Haskell
(9.10.1)
結果
TLE  
実行時間 -
コード長 2,535 bytes
コンパイル時間 2,819 ms
コンパイル使用メモリ 178,432 KB
実行使用メモリ 290,976 KB
最終ジャッジ日時 2024-06-26 02:30:59
合計ジャッジ時間 8,007 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other TLE * 1 -- * 27
権限があれば一括ダウンロードができます
コンパイルメッセージ
Loaded package environment from /home/judge/.ghc/x86_64-linux-9.8.2/environments/default
[1 of 2] Compiling Main             ( Main.hs, Main.o )

Main.hs:45:15: warning: [GHC-63394] [-Wx-partial]
    In the use of ‘head’
    (imported from Prelude, but defined in GHC.List):
    "This is a partial function, it throws an error on empty lists. Use pattern matching or Data.List.uncons instead. Consider refactoring to use Data.List.NonEmpty."
   |
45 |       let s = head (head vs)
   |               ^^^^

Main.hs:45:21: warning: [GHC-63394] [-Wx-partial]
    In the use of ‘head’
    (imported from Prelude, but defined in GHC.List):
    "This is a partial function, it throws an error on empty lists. Use pattern matching or Data.List.uncons instead. Consider refactoring to use Data.List.NonEmpty."
   |
45 |       let s = head (head vs)
   |                     ^^^^

Main.hs:80:17: warning: [GHC-63394] [-Wx-partial]
    In the use of ‘head’
    (imported from Prelude, but defined in GHC.List):
    "This is a partial function, it throws an error on empty lists. Use pattern matching or Data.List.uncons instead. Consider refactoring to use Data.List.NonEmpty."
   |
80 |   let n = read (head inputN) :: Int
   |                 ^^^^
[2 of 2] Linking a.out

ソースコード

diff # 

module Main where

-- https://yukicoder.me/problems/no/2356
split :: String -> [String]
split = go ("", [])
  where
    go :: (String, [String]) -> String -> [String]
    go (a, acc) "" = acc <> [a]
    go (a, acc) (' ' : xs) = go ("", acc <> [a]) xs
    go (a, acc) (x : xs) = go (a <> [x], acc) xs

prompt :: IO [String]
prompt = do split <$> getLine

promptN :: Int -> IO [[String]]
promptN n = go n []
  where
    go :: Int -> [[String]] -> IO [[String]]
    go 0 acc = return acc
    go n acc = do
      vs <- prompt
      go (n - 1) (acc <> [vs])

fixMod :: Int
fixMod = 998244353

pow :: Int -> Int -> Int -> Int
pow _ 0 _ = 1
pow x n m =
  let x' = x * x `mod` m
      n' = n `div` 2
      r = pow x' n' m
   in if odd n
        then (x * r) `mod` m
        else r

inv :: Int -> Int -> Int
inv x m = pow x (m - 2) m

extract :: [[String]] -> ([Char], [Int], Int)
extract = foldr f ([], [], 0)
  where
    f :: [String] -> ([Char], [Int], Int) -> ([Char], [Int], Int)
    f vs (ss, as, doorCount) =
      let s = head (head vs)
          a = read (last vs) :: Int
       in (s : ss, a : as, doorCount + a)

calcSeason :: Int -> [Char] -> [Int] -> (Int, Int, Int, Int)
calcSeason n ss as = go n ss as (0, 0, 0, 0)
  where
    go :: Int -> [Char] -> [Int] -> (Int, Int, Int, Int) -> (Int, Int, Int, Int)
    go 0 _ _ acc = acc
    go m (s : restS) (a : restA) (summer, fall, winter, spring) =
      let prd = pow (n - 1) a fixMod
          prd' = (prd + fixMod - pow (n - 2) a fixMod) `mod` fixMod
          prd'' = (prd' * inv (pow (n - 1) a fixMod) fixMod) `mod` fixMod
       in go
            (m - 1)
            restS
            restA
            ( case s of
                'U' -> ((summer + prd'') `mod` fixMod, fall, winter, spring)
                'F' -> (summer, (fall + prd'') `mod` fixMod, winter, spring)
                'W' -> (summer, fall, (winter + prd'') `mod` fixMod, spring)
                'P' -> (summer, fall, winter, spring + 1)
            )

solve :: Int -> Int -> Int -> Int -> Int -> Int -> Int
solve n doors summer fall winter spring =
  let a' = (1 * summer) `mod` fixMod
      a'' = (a' * fall) `mod` fixMod
      a''' = (a'' * winter) `mod` fixMod
      answer = (a''' * spring) `mod` fixMod
   in (answer * pow (n - 1) doors fixMod) `mod` fixMod

main :: IO ()
main = do
  inputN <- prompt
  let n = read (head inputN) :: Int
  vs <- promptN n
  let (ss, as, doorCount) = extract vs
      (summer, fall, winter, spring) = calcSeason n ss as
  print $ solve n doorCount summer fall winter spring
0