yukicoder 661, 662, 663, 664, 665

最近いろんな作業を Rust で書いてるんですが、Rust やってると Haskell が書きたくなってくる。664はまだです。

661 ハローキティはりんご3個分

import Control.Monad

main = getLine >> getContents >>= mapM_ (putStrLn . solve . read) .  lines

solve n
  | n `mod` 40 == 0 = "ikisugi"
  | n `mod` 8 == 0 = "iki"
  | n `mod` 10 == 0 = "sugi"
  | otherwise = show (n `div` 3)

662 スロットマシーン

import Data.List
import Data.Maybe
import Control.Monad

main = do
  dic <- replicateM 5 readKV
  (n1, a) <- readReel (map fst dic)
  (n2, b) <- readReel (map fst dic)
  (n3, c) <- readReel (map fst dic)

  -- [0] [1] [2] [3] [4]
  -- 000         0     0
  --     000      0   0
  --         000   0 0

  let cnt = map (\k -> let x = a!!k; y = b!!k; z = c!!k in 5*x*y*z) $ [0..4]
  let ans = sum $ zipWith (*) cnt (map snd dic)

  print $ fromIntegral ans / fromIntegral (n1 * n2 * n3)
  mapM_ print cnt


readKV :: IO (String, Int)
readKV = do
  [k, v] <- words <$> getLine
  return (k, read v)

readReel :: [String] -> IO (Int, [Int])
readReel dic = do
  n <- readLn :: IO Int
  x <- replicateM n (fromJust . flip elemIndex dic <$> getLine)
  return (n, [length . filter (==i) $ x | i <- [0..4]])

663 セルオートマトンの逆操作

リストモナドの bind を連鎖させれば全列挙できるけど、それだと指数時間になってしまう。bind 後に重複除去をすると計算量が落ちる（いまさらなんだけど、groupBy する前にソートしないと危険）。dp[i] から dp[i+1] が作れる系の DP はたいていこう書けると思う。

import Data.List
import Control.Monad

main = do 
  n <- readLn :: IO Int
  a <- replicateM n readLn :: IO [Int]
  print $ solve a

md :: Int
md = 10^9 + 7

type Stat = (Int, Int, Int, Int)

solve :: [Int] -> Int
solve a = let b = foldl' g [((i, j, i, j), 1) | i <- [0,1], j <- [0,1]] a
          in modsum . map snd . filter (\((h1, h2, x, y), _) -> h1 == x && h2 == y) $ b
  where
    g :: [(Stat, Int)] -> Int -> [(Stat, Int)]
    g acc a = compress (acc >>= f a)

    f :: Int -> (Stat, Int) -> [(Stat, Int)]
    f a ((h1, h2, x, y), cnt) = [((h1, h2, y, z), cnt) | z <- [0, 1], rule x y z == a]

compress :: [(Stat, Int)] -> [(Stat, Int)]
compress a = map f . groupBy (\x y -> fst x == fst y) $ a
  where
    f :: [(Stat, Int)] -> (Stat, Int)
    f x = let key = fst (head x); val = modsum . map snd $ x
          in (key, val)

modsum :: [Int] -> Int
modsum = foldl' (\x y -> (x + y) `mod` md) 0

rule :: Int -> Int -> Int -> Int
rule 0 0 0 = 0
rule 0 0 1 = 1
rule 0 1 0 = 1
rule 0 1 1 = 1
rule 1 0 0 = 0
rule 1 0 1 = 1
rule 1 1 0 = 1
rule 1 1 1 = 0

665 Bernoulli Bernoulli

F# だと演算子の優先順位がいい感じに自動生成されるんだけど、Haskell でその感覚でやったらハマった（どういう順序になってるんでしょう）。

import Data.List (foldl', scanl)

md :: Int
md = 10^9 + 7

-- sometimes ambigious
infixl 6 +++
(+++) :: Int -> Int -> Int
(+++) a b = (a + b) `mod` md

infixl 7 ***
(***) :: Int -> Int -> Int
(***) a b = a * b `mod` md

infixr 8 ^^^
(^^^) :: Int -> Int -> Int
(^^^) a b
  | b == 0 = 1
  | odd b = a *** a ^^^ (b - 1)
  | otherwise = (a *** a) ^^^ (b `div` 2)

modinv :: Int -> Int
modinv a = a ^^^ (md - 2)

infixl 7 ///
(///) :: Int -> Int -> Int
(///) a b = a *** modinv b

prodmod :: [Int] -> Int
prodmod = foldl' (***) 1

--         p/(x-0)                                 p/(x-1)
--  vvvvvvvvvvvvvvvvvvvv                    vvvvvvvvvvvvvvvvvvvv
--  (x-1)(x-2)...(x-n+1)            (x-0)   (x-2)(x-3)...(x-n+1)
-- --------------------- a[0] + ... ---------------------------- a[1] + ...
--  (0-1)(0-2)...(0-n+1)            (1-0)   (1-2)(1-3)...(1-n+1)
--                                  ^^^^^   ^^^^^^^^^^^^^^^^^^^^
--                                   [L]            [R]

lagrange :: [Int] -> Int -> Int
lagrange a x
  | x < n = a !! x
  | otherwise = let (sm, _, _, _) = foldl' f (0, 0, 1, q) a
                in sm
  where
    n = length a
    p = prodmod [(x - i) | i <- [0..n-1]]
    q = prodmod [-i | i <- [1..n-1]]
    f :: (Int, Int, Int, Int) -> Int -> (Int, Int, Int, Int)
    f (sm, k, l, r) ak = (nextSm, k + 1, nextL, nextR)
      where 
        nextSm = sm +++ ak *** p /// (l *** r *** (x-k))
        nextL = l *** (k + 1)
        nextR = r /// negate (n - 1 - k)

computeSmall :: Int -> [Int]
computeSmall k = scanl f 0 [1..k + 1]
  where
    f s x = s +++ (x ^^^ k)

solve :: Int -> Int -> Int
solve n k = lagrange (computeSmall k) (n `mod` md)

main = do
  [n, k] <- map read . words <$> getLine :: IO [Int]
  print $ solve n k