ABC250 A～E をHaskellで

A - Adjacent Squares

問題 ABC250A

シグネチャを決める。

abc250a :: Int -> Int   -- H,W
        -> Int -> Int   -- R,C
        -> Int

「角なら2」「角でない辺なら3」「角でも辺でもないなら4」とやると、$H=1$とか$W=1$とかその両方とかの落とし穴にはまる。
どうせたかだか4なので、全部作って調べて数える。

結果

abc250a h w r c = length
  [ ()
  | (dx,dy) <- [(1,0),(0,1),(-1,0),(0,-1)]
  , let c1 = c + dx, let r1 = r + dy
  , 1 <= r1, r1 <= h, 1 <= c1, c1 <= w
  ]

B - Enlarged Checker Board

問題 ABC250B

シグネチャを決める。

abc250b :: Int -> Int -> Int  -- N,A,B
        -> [String]           -- 出力

古風な、アスキーアート的な画面を作る問題。
白 . または黒 # の $B$ 文字の並びを $N$ 個並べることを $A$ 行繰り返すことを $N$ 回繰り返す。そのまま書くと

abc250b n a b =
  concat $ take n $ cycle $
    [replicate a $ concat $ take n $ cycle [replicate b '.', replicate b '#']
    ,replicate a $ concat $ take n $ cycle [replicate b '#', replicate b '.']]

concat $ take n $ cycle $ map (replicate K) [{-X-}, {-Y-}] という構造が2重に現れていることがわかるので、そこを括りだす。

結果

abc250b n a b = sub b x y
  where
    x = sub a '.' '#'
    y = sub a '#' '.'
    sub k x y = concat $ take n $ cycle $ map (replicate k) [x, y]

abc250b n a b = sub b [sub a ".#", sub a "#."]
  where
    sub k = concat . take n . cycle . map (replicate k)

C - Adjacent Swaps

問題 ABC250C

シグネチャを決める。

abc250c :: Int -> Int   -- N,Q
        -> [Int]        -- xi
        -> [Int]        -- 答えai

位置 $i$ のボールに書いてある整数を配列 $A[0 \leq i \leq N-1]$ で管理すると、整数 $x_i$ の書いてあるボールの現在位置を知るのに $A$ を線形探索する羽目になる。
整数 $x$ の書いてあるボールの位置を配列 $B[1 \leq i \leq N]$ で管理すると、位置 $B[x_i]+1$ にある、入れ替えるべきボールに書いてある整数を知るのに $B$ を線形探索する羽目になる。
しかしこの配列が両方ともあれば、一方を更新するのに必要な情報はもう一方から $O(1)$ で得られる。

結果

全体をSTモナドアクションに変更した。単なる関数呼び出しの代わりに runST で起動する。

import Control.Monad
import Control.Monad.ST
import Data.Array.ST

abc250c :: Int -> Int -> [Int] -> ST s [Int]
abc250c n q xs = do
  i2x <- newArray_ (1, n) :: ST s (STUArray s Int Int)
  x2i <- newArray_ (1, n) :: ST s (STUArray s Int Int)
  forM_ [1..n] (\i -> do
    writeArray i2x i i
    writeArray x2i i i
    )
  forM_ xs (\x -> do
    i <- readArray x2i x
    let j = if i == n then pred i else succ i
    y <- readArray i2x j
    swap x2i i j
    swap i2x x y
    )
  forM [1..n] (readArray i2x)

swap a p q = do
  s <- readArray a p
  t <- readArray a q
  writeArray a p t
  writeArray a q s

素直にIOモナド、IOArrayを使い、$x_i$ をひとつ読むごとに処理する形の方が節約になる。

D - 250-like Number

問題 ABC250D

問題文も例も紛らわしく、$q$ が素数なのかどうか不明瞭であるが、結論としては $q$ も素数限定のようだ。

シグネチャを決める。

abc250d :: Int -> Int

$k = p \times q^3, p < q$ とは、$k$ がそのように素因数分解されるということなので、異なる素数対 $r,s$ によって $k=r \times s^3$ となることはありえない。
そこで、素数 $p$ について、それより大きな素数 $q$ で $p \times q^3 \leq n$ となるものの個数を数えて総和をとる。（$k$ を作って重複を除いて数える必要がない、ということである。）
$q$ の候補は、$q^3 \leq n/p$ を満たす範囲と特定できる。
そのような $q$ の候補が0個になるとき、$p$ の走査上限に達したとわかる。

$q$の候補の上限を探すのにナイーブに $p q^3 \leq n$ とやると、$n \leq 10^{18}$ と大きいことから64ビット整数をオーバーフローするという罠がある。

結果

import Data.List

abc250d :: Int -> Int
abc250d n = sum
  [ length q3s
  | (p, q3s) <- takeWhile (not . null . snd) $ map (pq3gen n) $ tails primes
  ]

pq3gen k (p:qs) = (p, takeWhile (div k p >=) $ map (^ 3) qs)

$q$の候補の上限を二分探索するべきだったかもしれないが、線形探索で足りた。

E - Prefix Equality

問題 ABC250E

シグネチャを決める。

abc250e :: Int              -- N
        -> [Int] -> [Int]   -- ai, bi
        -> Int              -- Q
        -> [(Int,Int)]      -- xi, yi
        -> [Bool]           -- 答え

$A$ の先頭 $x$ 項に含まれる値の集合を $A_x$ とする。

$b_j$ のそれぞれの値について、最小の出現位置がどこなのかを疎な表にする。出現しないとき $N+1$ またはより大きな値とする。
$\textrm{bminpos}(b) = \min \; \{ j \;|\; b_j = b \}$
$A$ のそれぞれの位置 $i$ の値 $a_i$ について、$B$ で出現する最小位置がどこなのかを、上の表から調べる。
$\textrm{xminy}(i) = \textrm{bminpos}(a_i)$

$A_x$ の全ての要素が $B_y$ に含まれるためには、$a_1, \dots, a_x$ が $\textrm{bminpos}(a_i) \leq y$ である必要がある。言い換えると $\forall i \; (1 \leq i \leq x) \; \textrm{xminy}(i) \leq y$ である。より手前の条件が重なることから、$\max$ でこれを累積する。
$\textrm{xminyAcc}(x) = \max_{1 \leq i \leq x} \textrm{xminy}(i)$
すると上の条件は $\textrm{xminyAcc}(x) \leq y$ だけで表せる。

$A$ と $B$ を入れ替えた版も作って、$\textrm{yminAcc}(y) \leq x$ も満たされるなら、$A_x = B_y$ であるといえる。

結果

import qualified Data.IntMap as IM
import Data.Array

abc250e :: Int -> [Int] -> [Int] -> Int -> [(Int,Int)] -> [Bool]
abc250e n as bs q xys = map cond xys
  where
    cond (x, y) = xminyAcc ! x <= y && yminxAcc ! y <= x
    bminpos = IM.fromListWith min $ zip bs [1..]
    aminpos = IM.fromListWith min $ zip as [1..]
    xminy = map (\a -> IM.findWithDefault maxBound a bminpos) as
    yminx = map (\b -> IM.findWithDefault maxBound b aminpos) bs
    xminyAcc = listArray (1,n) $ scanl1 max xminy
    yminxAcc = listArray (1,n) $ scanl1 max yminx

蛇足だが、$a,b$ の変域は$10^9$までと大きく疎なので IntMap を使い、$x,y$ の変域は$2 \times 10^5$までで隙間なく値が存在するので Array を用いる、と使い分けている。