ABC281 A～E をHaskellで

A - Count Down

問題 ABC281A

$N$から0までのカウントダウンはレンジで生成できる。

結果

main = do
  n <- readLn
  mapM_ print [n, pred n .. 0]

B - Sandwich Number

問題 ABC281B

シグネチャを決める。

abc281b :: String  -- S
        -> Bool    -- 答え

条件を整理すると、$S$は8文字であること、先頭と末尾は英大文字であること、中6文字は数字で、先頭は0ではないこと、が確認できればよいとわかる。

結果

import Data.Char

abc281b :: String -> Bool
abc281b (x:xs) =
    isUpper x && length xs == 7 && isUpper (last xs) &&
    all isDigit (init xs) && head xs /= '0'

C - Circular Playlist

問題 ABC281C

シグネチャを決める。

abc281c :: Int        -- N
        -- Int        -- T
        -- [Int]      -- Ai
        -- (Int,Int)  -- 答え 番号/経過時間

ループするので、$T \bmod \sum A_i$ が先頭から何秒なのかを考えればよい。
本番中はうっかり IntSet.lookupLT を使ってしまったが、検索するのは一度きりなので線形探索で充分である。

結果

abc281c :: Int -> Int -> [Int] -> (Int,Int)
abc281c n t as = (length ts, t1 - last ts)
  where
    t1 = mod t (last accs)
    accs = scanl' (+) 0 as
    ts = takeWhile (t >) accs

D - Max Multiple

問題 ABC281D

シグネチャを決める。

abc281d :: Int    -- N
        -> Int    -- K
        -> Int    -- D
        -> [Int]  -- Ai
        -> Int    -- 答え

3次元配列 $arr[0 \leq I \leq N][0 \leq J \leq \min(I,K)][0 \leq r \leq D]$ は、$A_i$ の前から $0 \leq I \leq N$ 個のうち $0 \leq J \leq \min(I,K)$ 個を使って、和 $s$ の $s \bmod D = r \; (0 \leq r < D)$ となるような最大の $s$ またはそのような組み合わせが存在しないとき $-1$ を持つとする。
このとき問題の答えは $arr[N][K][0]$ にある。
$arr[I][J][r]$ は、一つ少ない候補で同じことをする $arr[I-1][J][r]$ と、$A_I$ を使い、あと$I-1$個から$J-1$ 個で $(r - A_I) \bmod D$ を作る最大の結果（もしあれば）に $A_I$ を足したものの大きい方である。
$arr[I][0][r]$ は、$r = 0$ について$0$、それ以外について$-1$である。
$arr[0][J][r]$ は、$J = 0, r = 0$ について$0$（上と同じ）、それ以外について$-1$である。

漸化関係をみると、つまり$I$ の次元について直前のものしか使わないので、3次元配列全体を作らずとも、2次元配列を順々に作れば求められる。
この問題は $N \leq 100$ と規模が抑えられているので、3次元配列でも実行できる。

結果

3次元配列でする版

import Data.Array

abc281d :: Int -> Int -> Int -> [Int] -> Int
abc281d n k d as = arr ! (n,k,0)
  where
    a = listArray (1,n) as
    bnds = ((0,0,0),(n,k,pred d))
    arr = listArray bnds $ map fun $ range bnds
    fun (_,0,0) = 0
    fun (_,0,_) = -1
    fun (0,_,_) = -1
    fun (i,j,r)
      | cand2s == -1 = cand1
      | otherwise    = max cand1 cand2
      where
        ai = a ! i
        cand1  = arr ! (pred i,j,r)
        cand2s = arr ! (pred i, pred j, mod (r - ai) d)
        cand2  = ai + cand2s

$I=0,1,\dots,N$ の2次元配列を順に構成する版

import Data.List
import Data.Array.Unboxed

type UA = UArray (Int,Int) Int
 
abc281d :: Int -> Int -> Int -> [Int] -> Int
abc281d n k d as = arrN ! (k,0)
  where
    bnds = ((0,0),(k,pred d))
    arr0, arrN :: UA
    arr0 = accumArray (flip const) (-1) bnds [((0,0),0)]
    arrN = foldl' {-'-} step arr0 $ zip [1..] as
    step :: UA -> (Int,Int) -> UA
    step arrI1 (i, ai) = listArray bnds $ map funI $ range bnds
      where
        funI (0,0) = 0
        funI (0,_) = -1
        funI (j,r)
          | cand2s == -1 = cand1
          | otherwise    = max cand1 cand2
          where
            cand1  = arrI1 ! (j,r)
            cand2s = arrI1 ! (pred j, mod (r - ai) d)
            cand2  = ai + cand2s

後者の方が、正格な foldl' と UArray により $I=0,1,\dots$ の順に配列の内容を確定させることができ、高速かつ省メモリである。（269ms vs 55ms, 121MB vs 6.6MB）

E - Least Elements

問題 ABC281E

シグネチャを決める。

abc281d :: Int    -- N
        -> Int    -- M
        -> Int    -- K
        -> [Int]  -- Ai
        -> [Int]  -- 答え

$A_i$ の $A_i$から始まる $M$ 要素について、小さい方から $K$ 個の和を、全ての $i$ について求めよ、という要求。

連続する$M$要素の区間が右にズレていくところから、尺取り法的な全体構造を想起する。
先頭$M$要素 $A_1, \dots, A_M$ については、ストレートに、整列して小さい方から $K$ 個と残りに分けることができる。前者の和が最初の答えである。

以降、$M$要素の範囲を順に右にずらして、変動だけを追跡する。
現在の$M$個の範囲で、小さい方から$K$個の集団を$G_L$、残りを$G_H$、$G_L$の要素の和を$s$とする。
範囲を右にずらすときに、集団から抜ける値を$A_L$、加わる値を$A_R$とする。
$A_L$が$G_L, G_H$ のどちらから抜けるか、$A_R$が$G_L,G_H$のどちらに加わるかの4通りで場合分けして、$s, G_L, G_H$ を正しく更新する。

なお、$K=M$のとき、$G_H$が常に空となり、範囲の全ての要素が和に加わるので、この場合は別に分けて書く。

結果

集団$G_L, G_H$は、同じ値が複数個同時に存在しうるので、単なる集合ではなく多重集合で実現する必要がある。また、片方の最小値だけが取り出せるヒープでは、「最大値を得る、取り除く」という操作を実現するのは符号反転で何とかなっても、「任意の要素を取り除く」という操作が効率的にできない。多重集合そのものの表現が必要である。

（自分はこれを単なる集合で誤っているのに気づかずに終了してしまった。しかしそれをチェックできるテストケースが1つだけだったのも少し意外。）

import qualified Data.IntMap as IM

abc281e :: Int -> Int -> Int -> [Int] -> [Int]
-- 区間の全ての要素が和に加わる K=M の場合
abc281e n m k as
  | m == k = scanl step sum0 $ zip as as2
  where
    (as1,as2) = splitAt m as
    sum0 = sum as1
    step sum0 (aL,aH) = sum0 - aL + aH
-- 一般の場合
abc281e n m k as = (sum0 :) $ snd $ mapAccumL step (sum0, fromListMIS as1L, fromListMIS as1H) $ zip as as2
  where
    (as1, as2) = splitAt m as
    (as1L, as1H) = splitAt k $ sort as1
    sum0 = sum as1L
    step (sum0, setL, setH) (aL, aR)
      | aL <= findMaxMIS setL = step2L (sum0 - aL) (deleteMIS aL setL) setH
      | otherwise             = step2H sum0 setL (deleteMIS aL setH)
      where
        step2L sum0 setL setH -- Lがひとつ少ない
          | aR <= findMinMIS setH = ((sum0H, insertMIS aR setL, setH), sum0H)
          | otherwise = ((sum0v, insertMIS v setL, insertMIS aR setH1), sum0v)
          where
            sum0H = sum0 + aR
            sum0v = sum0 + v
            v = findMinMIS setH
            setH1 = deleteMIS v setH
        step2H sum0 setL setH -- Hがひとつ少ない
          | aR >= findMaxMIS setL = ((sum0, setL, insertMIS aR setH), sum0)
          | otherwise = ((sum0vH, insertMIS aR setL1, insertMIS v setH), sum0vH)
          where
            v = findMaxMIS setL
            setL1 = deleteMIS v setL
            sum0vH = sum0 - v + aR

-- IntMapによる多重集合の実装
-- @gotoki_no_joe
type MultiIntSet = IM.IntMap Int
fromListMIS xs = IM.fromListWith (+) [(x, 1) | x <- xs]
insertMIS x ms = IM.insertWith (+) x 1 ms
deleteMIS x ms = IM.update dec x ms
  where
    dec 1 = Nothing
    dec n = Just $ pred n
findMaxMIS = fst . IM.findMax
findMinMIS = fst . IM.findMin