ABC260 A～E をHaskellで

A - A Unique Letter

問題 ABC260A

シグネチャを決める。

abc260a :: String  -- S
        -> String  -- 答え

一度だけ現れる文字が唯一に定まる、という優しい前提がもしあれば、ある二つが等しいときもう一つは必ず仲間外れ、というやり方でできる。

-- WA
abc260a [a,b,c]
  | a == b = [c]
  | b == c = [a]
  | c == a = [b]

しかし今回はそうではないので、条件を満たすか、つまり他の2つと異なるものかどうかを丁寧に確認する必要がある。

結果

abc260a [a,b,c]
  | cond a b c = [a]
  | cond b c a = [b]
  | cond c a b = [c]
  | otherwise = "-1"

cond a b c = a /= b && a /= c

B - Better Students Are Needed!

問題 ABC260b

シグネチャを決める。

abc260b :: Int                -- N
        -> Int -> Int -> Int  -- X,Y,Z
        -> [Int] -> [Int]     -- Ai, Bi
        -> [Int]

要求どおりに粛々と合格者を選び出していく。(番号, 数学の点, 英語の点) というタプルについて、様々な比較条件でソートをかける。
降順でソートするにはよくflip compareを用いるが、値をnegateすることでもできる。
「同点のとき番号の小さい方を優先」は、タプルの順序が辞書式と定義されていることを利用して実現する。

結果

import Data.List

abc260b n x y z as bs = sort $ map (\(i,_,_) -> i) $ p1 ++ p2 ++ p3
  where
    iabs0 = zip3 [1..] as bs
    (p1, iabs1) = splitAt x $ sortOn (\(i,a,_) -> (-a,i)) iabs0
    (p2, iabs2) = splitAt y $ sortOn (\(i,_,b) -> (-b,i)) iabs1
    p3          = take z  $ sortOn (\(i,a,b) -> (-a-b,i)) iabs2

C - Changing Jewels

問題 ABC260C

シグネチャを決める。

abc260c :: Int -> Int -> Int  -- N,X,Y
        -> Int                -- 答え

(レベル $n$ の赤い宝石 $r$ 個, レベル $n$ の青い宝石 $b$ 個) という状態から、レベル $n$ の赤い宝石を全て変換すると、(レベル $n-1$ の赤い宝石 $r$ 個, レベル $n$ の青い宝石 $b + rX$ 個) という状態になる。（赤い宝石のレベルに注意）

red  (r, b) = (r, b + r * x)

同様に、(レベル $n-1$ の赤い宝石 $r$ 個, レベル $n$ の青い宝石 $b$ 個) という状態から、レベル $n$ の青い宝石を全て変換すると、(レベル $n-1$ の赤い宝石 $r + b$ 個, レベル $n-1$ の青い宝石 $bY$ 個) という状態になる。

blue (r, b) = (r + b, b * y)

この変換はよく見ると、実行の順序は結果に影響せず、レベル $1$ の青い宝石を最大にするには、この2つの手順を繰り返して全てをレベル $1$ にすることしかできない。

結果

abc260c n x y = snd $ (!! pred n) $ iterate (blue . red) (1, 0)
  where
    red  (r, b) = (r, b + r * x)
    blue (r, b) = (r + b, b * y)

別解

blue . red の合成関数を手計算して、レベル $n$ の (赤い宝石 $r$ 個, 青い宝石 $b$ 個) は、レベル $n-1$ の (赤い宝石 $r(X+1) + b$ 個, 青い宝石 $Y(rX + b)$ 個) になるという関係を求め、レベルとともに追跡することもできる。

abc260c n x y = loop n 1 0
  where
    loop 1 _ b = b
    loop n r b = loop (pred n) (r * succ x + b) (y * (r * x + b))

D - Draw Your Cards

問題 ABC260D

シグネチャを決める。

abc260d :: Int -> Int   -- N, K
        -> [Int]        -- Pi
        -> [Int]        -- 答え

場に表になっているカードの山を、Data.IntMap で模倣する。一番上のカードの番号をキー、その山の枚数とカードのリストを値とする。
カード $X$ を重ねるべき山は lookupGT で選ぶことができる。
カード $K$ をターン $T$ で食べたという情報をリストにして集める。食べずに終わったカードの処理もあるし、sort で整列させるより配列に取り込む方が早く構築できる。

結果

import qualified Data.IntMap as IM
import Data.Array

abc260d :: Int -> Int -> [Int] -> [Int]
abc260d n k ps = elems $ accumArray (flip const) (-1) (1, n) $ concat ktss
  where
    (_, ktss) = mapAccumL step IM.empty $ zip ps [1..]
    step m (pi, i)
      | cnt1 == k = (m1, zip qs1 $ repeat i)
      | otherwise = (IM.insert pi (cnt1, pi:qs1) m1, [])
      where
        (m1, cnt1, qs1) =
          case IM.lookupGT pi m of
            Nothing             -> (            m, 1       , [])
            Just (q, (cnt, qs)) -> (IM.delete q m, succ cnt, qs)

各ステップの処理において、$K=1$ という場合もあるので、カードを重ねてみる処理と、カードが $K$ 枚重なったかを確認する処理の段階を分けると安全に実装できる。

E - At Least One

問題 ABC260E

シグネチャを決める。

abc260e :: Int -> Int   -- N, M
        -> [(Int,Int)]  -- (Ai, Bi)のペアのリスト abs
        -> [Int]        -- 答え

（自分も解説を見るまでまるで気づかず見当違いの方法を模索していたが、そんなことは棚に上げて）
$1$ から $M$ までの様々な区間について、条件を満たすものを見つけたい。区間を $[x,y]$ のように表し、ここで、区間 $[x,y]$ が条件を満たすとき、それを含む区間 $[w,z] \; (w < x, y < z)$ はやはり条件を満たすという性質があるとき、尺取り法が使える（そうな）。

尺を取るとき「$A_i$ または $B_i$ が現在の区間 $[lb,ub]$ に含まれるような $i$ の個数」を追跡する。

• これが $N$ のとき、区間は条件を満たしている。
• 上限を右に1つ広げたとき、$A_i, B_i$ のいずれかが $ub + 1$で、もう一方は区間の外であるような $i$ の個数だけ、含まれるような $i$ の個数は増える。
• 下限を右に1つ狭めたとき、$A_i, B_i$ のいずれかが $lb$で、もう一方は区間の外であるような $i$ の個数だけ、含まれるような $i$ の個数は減る。

下限を1ずつ進める外側のループの中で、条件を満たす上限を見つける内側のループが回る、という二重ループの構造にする代わりに、下限または上限のいずれかを1ずつ進めるステップを、続く限り続ける構成で実装する。状態を (下限, 上限, $i$ の個数) とし、区間が条件を満たすなら下限を進め、満たさないなら上限を進める。上限が $M$ を踏み越えたらそこで止める。

lucs = takeWhile (\(_,ub,_) -> ub <= m) $ iterate step (1,0,0)
step (lb, ub, cnt)
  | cnt == m  = (lb1, ub, cnt1)
  | otherwise = (lb, ub1, cnt2)
  where
    lb1 = succ lb
    ub1 = succ ub
    cnt1 = cnt - {- 下限をずらすことで外れるiの個数 -}
    cnt2 = cnt + {- 上限をずらすことで入るiの個数 -}

$i$ の個数の増減を数えるには、$A_i$ または $B_i$ の値をキーとし、ペアの相手側 $B_i$ または $A_i$ を値とする配列（ただし値は複数になりうるのでリストにする）を作っておく。
これを今から加える $ub+1$ または外す $lb$ でアクセスすると、区間の外にあるかどうか確認するべきペアのもう一方のリストが得られる。

ab2ba = accumArray (flip (:)) [] (1,m) [p | (a,b) <- abs, p <- [(a,b), (b,a)]]

cnt1 = cnt - length [() | x <- ab2ba ! lb , x < lb1 || ub  < x]
cnt2 = cnt + length [() | x <- ab2ba ! ub1, x < lb  || ub1 < x]

尺取りが終わったら、lucsの中でカウントが $N$ に等しいもの $(lb,ub,N)$ についてそれぞれ、区間 $[lb,ub]$（長さ $ub-lb+1$）、区間 $[lb,ub+1]$（長さ $+1$）、…、区間 $[lb,M]$（長さ $M-lb+1$）が一つある、という個数を数えれば答えが得られる。

elems $ accumArray (+) 0 (1,m) $
[(x - lb + 1, 1) | (lb,ub,cnt) <- lucs, cnt == N, x <- [ub .. M]]

ここでこのナイーブな方法を用いると、広い範囲の二重ループで大変なことになる。そこで累積和を利用する。$(lb,ub,N)$ があることで、$f(\cdot)$ は $f(ub-lb+1) - f(ub-lb) = 1$ の位置で1増加し、$f(M-lb+1) - f(M-lb+2) = 1$ の位置で1減少する。この $f$ の差分を配列に累積し、最後に左から足しこむことで答えを得る。

take m $ scanl1 (+) $
elems $ accumArray (+) 0 (1, succ m) $
[p | (lb,ub,cnt) <- lucs, cnt == N, p <- [(ub-lb+1, 1), (m-lb+2, -1)]]

結果

import Data.Array

abc260e :: Int -> Int -> [(Int,Int)] -> [Int]
abc260e n m abs = take m $ scanl1 (+) $ elems df
  where
    ab2ba = accumArray (flip (:)) [] (1,m) [p | (a,b) <- abs, p <- [(a,b), (b,a)]]
    lucs = takeWhile (\(_,ub,_) -> ub <= m) $ iterate step (1,0,0)
    step (lb, ub, cnt)
      | cnt == m  = (lb1, ub, cnt1)
      | otherwise = (lb, ub1, cnt2)
      where
        lb1 = succ lb
        ub1 = succ ub
        cnt1 = cnt - length [() | x <- ab2ba ! lb , x < lb1 || ub  < x]
        cnt2 = cnt + length [() | x <- ab2ba ! ub1, x < lb  || ub1 < x]
    df = accumArray (+) 0 (1, succ m) $
         [p | (lb,ub,cnt) <- lucs, cnt == N, p <- [(ub-lb+1, 1), (m-lb+2, -1)]]