ABC291 A～F をHaskellで

(2023/2/27) Eの別解を追記。

A - camel Case

問題 ABC291A

先頭から、小文字である間（大文字が出現する手前まで）を取り出した個数+1が答え。

結果

import Data.Char

main = getLine >>= print . succ . length . takeWhile isLower

B - Trimmed Mean

問題 ABC291B

シグネチャを決める。

abc291b :: Int       -- N
        -> [Double]  -- Xi
        -> Double    -- 答え

まずソートしてしまえば、先頭N個を除き、さらに、末尾N個を除くためには、前3N個を取り出せばいい。

結果

import Data.List

abc291b :: Int -> [Double] -> Double
abc291b n xs = (sum $ take n3 $ drop n $ sort xs) / fromIntegral n3
  where
    n3 = n * 3

C - LRUD Instructions 2

問題 ABC291C

シグネチャを決める。

abc291c :: Int     -- N
        -> String  -- S
        -> Bool    -- 答え

出発点を適当に決め、移動するたびに座標を更新しつつ、過去に移動したことのある座標の集合を確認して、二度目なら True を返す。そうならなくて S が尽きたら False を返す。

結果

打ち切りありの計算なので、再帰関数で書き捨ててしまう。

import qualified Data.Set as S

abc291c :: Int -> String -> Bool
abc291c n s = loop (0,0) S.empty s

loop xy xys _ | S.member xy xys = True
loop _ _ [] = False
loop xy xys (c:s) = loop (step xy c) (S.insert xy xys) s

step (x,y) 'U' = (x, succ y)
step (x,y) 'D' = (x, pred y)
step (x,y) 'L' = (pred x, y)
step (x,y) 'R' = (succ x, y)

「普通のプログラミング言語」の先入観があるとついつい、関数定義のガード式で、パターンで受けた場合全てを漏れなく配分しないといけない、と思い込んでしまうが、実際には漏れがあると次の等式に進んでくれるので、こんな風に書ける。loop の定義の3行は

• xy が既出なら、重複
• $S$ が尽きたなら、重複なし
• $S$ がまだあるなら、現在位置を既出xysに追加し、次の移動を行い、繰り返し

という、自然な動作の説明とよく対応している。

網羅的でないガード節を敬遠すると、次のようにもコード化できるが、自分は、上が最も直観的に書けていると思っている。

-- S.memberのチェックが重複
loop xs xys (c:s)
  | S.member xs xys = True
  | otherwise       = loop (step xy c) (S.insert xy xys) s
loop xy xys [] = S.member xs xys -- oops.

-- 今回は結果がBoolなのでショートカット演算子でこう書けるけど、やはりmemberは重複
loop xy xys (c:s) = S.member xs xys || loop (step xy c) (S.insert xy xys) s
loop xy xys []    = S.member xs xys -- oops.

-- 衝突のチェックをするタイミングをずらす
-- 呼び出し側も、スタート地点を入れた S.singleton (0,0) を初期値にする
loop xy xys (c:s) = S.member xy1 xys || loop xy1 (S.insert xy1 xys) c
  where xy1 = step xy c
loop _ _ [] = False

-- 二段階に計算のステップを分ける
part1 xy xys s = S.member xy xys || part2 xy (S.insert xy xys) s
part2 xy xys (c:s) = part1 (step xy c) xys s
part2 _  _   []    = False

D - Flip Cards

問題 ABC291D

シグネチャを決める。$A_i, B_i$ はタプルにすべきだが、無精する。

abc291d :: Int     -- N
        -> [[Int]] -- Ai, Bi
        -> Int     -- 答え

ごく簡単なDP。
直前のカードを表のままにした（$A_i$が見えている）場合と裏返した（$B_i$が見えている）場合とで、禁止されていない並べ方の場合の数（それぞれ $X_i, Y_i$ とする）を集計していく。
1枚めのカードは表でも裏でもいいので場合は1通りずつ $(X_1=1, Y_1=1)$ である。
以降、

• $A_i \neq A_{i+1}$ ならばカード$i+1$は表にできるので、$X_{i+1}$ には $X_i$ を足しこむ
• $B_i \neq A_{i+1}$ ならばカード$i+1$は表にできるので、$X_{i+1}$ には $Y_i$ を足しこむ
• $A_i \neq B_{i+1}$ ならばカード$i+1$は裏にできるので、$Y_{i+1}$ には $X_i$ を足しこむ
• $B_i \neq B_{i+1}$ ならばカード$i+1$は裏にできるので、$Y_{i+1}$ には $Y_i$ を足しこむ
  で増やしていく。

結果

import Control.Monad
import qualified Data.ByteString.Char8 as BS
import Data.Char
import Data.List

add x y = mod (x + y) 998244353

abc291d :: Int -> [[Int]] -> Int
abc291d n ((a:b:_):abs) = post $ foldl step (a,1,b,1) abs

post (_,x,_,y) = add x y

step (a,x,b,y) (a1:b1:_) =
  ( a1, add (if a /= a1 then x else 0) (if b /= a1 then y else 0)
  , b1, add (if a /= b1 then x else 0) (if b /= b1 then y else 0)
  )

E - Find Permutation

問題 ABC291E

シグネチャを決める。$X_i, Y_i$ はタプルにすべきだが、無精する。
結果は空リストなら No の意味とする。

abc291e :: Int     -- N
        -> [[Int]] -- Xi, Yi
        -> [Int]   -- 答え

「より大きい」という有向辺からなるグラフが与えられて、全てのノードを一本道で辿れるか、辿れるなら出発点からの距離を各ノードについて順に言え、ということ。

普通の幅優先探索は $O(N)$ で完了できるが、今回求めたいのは最短距離ではないので、先走って短い距離を付けてしまったノードの距離を付けなおす、をしていると $O(N^2)$ かかってしまう。

「入力に矛盾しない$A$が存在する」の一文が重要。
つまり、$A_1 < A_2, A_2 < A_3, A_3 < A_1$ のような矛盾する情報が与えられることはない。
あるノード $A_j$ について、これより小さいとわかっている全ての $A_i$ の順位を求め、その最大値+1を $A_j$ の順位とする、より小さいが一人もいないノードは順位1とする、という lazy DP で順位をつけることができてしまう。（矛盾する情報があると、lazy DPの参照が循環して実行時エラーとなる。）

ここで情報が不足する場合は、順位の重複が起きるので、そのようなことがなければ成功と判断できる。

結果

import Data.Array

abc291e :: Int -> Int -> [[Int]] -> [Int]
abc291e n m xys
  | cond      = elems arr1
  | otherwise = []
  where
-- いつものようにグラフを作る 入り辺の方を登録する
    g = accumArray (flip (:)) [] (1,n) [(y,x) | (x:y:_) <- xys]
-- 集めるDP
    arr1 :: Array Int Int
    arr1 = listArray (1,n) $ map f [1..n]
    f :: Int -> Int
    f k = case g ! k of
-- 「より小さい」が一人もいないなら順位1
            [] -> 1
-- 「より小さい」ノードの順位の最大値+1
            xs -> succ $ maximum $ map (arr1 !) xs
-- 各順位のノードが何人になったか数える
    arr2 = accumArray (+) 0 (1,n) [(cnt, 1) |(a, cnt) <- assocs arr1]
-- 全て一人ずつなら成功、さもなくば失敗
    cond = all (1 ==) $ elems arr2

追記：標準的グラフアルゴリズムによる解法

この問題の内容は、グラフ理論の用語で「トポロジカルソートが可能なグラフが与えられる。そこにハミルトン（開）路があるならそれを見つけよ」と説明できるらしい。

Wikipediaに、トポロジカルソートを行うアルゴリズムが二つ紹介されている。グラフの入辺数が0なノードをグラフから取り除くことを繰り返す方法と、深さ優先探索で一番底まで降りることで逆順に発見する方法と。

前者は、グラフからノードを除いたときに、そこから出ている辺を取り除くところが、immutableな計算と相性がよくなさそうである。ただし、この方法を用いた場合、トポロジカルソートを完成させる前に、取り除く候補の入辺数が0なノードが、同時に複数存在したら即失敗と判定できるため、この問題には適している。

後者は非常に簡潔だが、ノードが到達済みであるというフラグ配列がmutable arrayを前提としている。こちらはトポロジカルソートを完成させてから、出てきた順序の隣接する要素間全てが、入力で大小関係を規定されていることが、問題の要求する条件になる。

STArrayを用いて後者のアルゴリズムでトポロジカルソートを行う別解を示す。

abc291e :: Int -> Int -> [[Int]] -> [Int]
abc291e n m xys
  | cond      = elems $ array (1,n) $ zip topsort [1..]
  | otherwise = []
  where
    g = accumArray (flip (:)) [] (1,n) [(x,y) | (x:y:_) <- xys]
    topsort = runST $ action n
    cond = and $ zipWith check topsort (tail topsort)
    check x y = elem y (g ! x)
 
    action :: Int -> ST s [Int]
    action n = do
      visited <- newArray (1,n) False
      foldM (visit visited) [] [1..n]
 
    visit :: STArray s Int Bool -> [Int] -> Int -> ST s [Int]
    visit v acc k = do
      x <- readArray v k
      if x then return acc else do
        writeArray v k True
        acc1 <- foldM (visit v) acc (g ! k)
        return (k:acc1)

F - Teleporter and Closed off

問題 ABC291E

シグネチャを決める。
データ量が多いこと、ランダムアクセスをしたい場面もありそうなことから、$S$はByteStringで扱う。

abc291f :: Int          -- N
        -> Int          -- M
        -> [ByteString] -- Si
        -> [Int]        -- 答え

普通に最短経路を調べて、その経路が踏んでいない都市についてはこれが答え、さもなくば…と考えてもラチがあかない。

都市は大きく$1$から$N$の向きに逆行することなく進むことしかできず、飛び幅も$1$から$M$までと制限されている。
ある都市$k$を通らない最短距離は、都市$k$より手前の都市$i$ $(k-M < i < k)$と、$k$より先の都市 $j$ $(k < j < k+M)$ で、$i$から直接 $k$ に行くことができるもの（つまり$k$を飛び越せるような都市 $i,j$ の組）について、「都市$1$から都市$i$の距離 $+$ 都市$j$から都市$N$の距離 $+1$」の最小値で得られる。
都市 $k$ を飛び越せる組が全くないときは、できない、と判断される。

それぞれの最短距離は今回の問題Eと同様にして求められる。（普通に幅優先探索で求めてもいい。）全てのノードが到達可能だいう保証はないことに注意。

組み立て

更新は不要だがランダムアクセスは重要な、配列を多用する。

import Data.Array

maxBoundは何かを足してしまうと負の値に回り込むので、そうはならない「とても大きな値」を用意しておく。

tooBig :: Int
tooBig = div maxBound 2

はじめ。

abc291f :: Int -> Int -> [BS.ByteString] -> [Int]
abc291f n m ss = answer
  where

与えられる $S_i$ は、都市 $i$ から行ける都市を列挙しているので、逆に辿ることで、都市 $N$ からの逆方向距離を数えるのは楽にできる。

-- distance reverse (backwards)
    dR = listArray (1,n) $ zipWith fR [1..n] ss
    fR i _ | i == n = 0
    fR i s = succ $ minimum $ tooBig :
             [dR ! j | (j, '1') <- zip [i+1..i+m] $ BS.unpack s]

一方、都市$1$からの順方向距離は、都市 $j$ に行ける都市の距離一覧を得るには、手前 $M$ 都市の $S$ の該当する位置を見ないとわからないので少し煩雑になる。

まず、各都市について、その手前$M$都市の$S$を、近い順（つまり逆順）に並べたリストを用意する。
不足分はオール '0' のダミーで埋める。

    sofs = transpose $ take m $ tail $ iterate (all0 :) ss
    all0 = BS.replicate m '0'

都市$j$に対するsofs の要素は順に、（0から数えて）第$d$要素の都市 $j-d-1$ の $S$ なので、その $d$ 文字目が '1' なら、この都市から $j$ に直接移動できる。

-- distance forward
    dF = listArray (1,n) $ zipWith fF [1..n] sofs
    fF 1 _ = 0
    fF j ss = succ $ minimum $ tooBig :
              [dF ! (pred j - d) | (d, s) <- zip [0..] ss, BS.index s d == '1']

細工は流々。いよいよ仕上げ。

全ての、都市$i$から$j$へのテレポート経路に関して、その間にあり飛び越される都市 $k$ $(i < k < j)$ について、$k$ を飛ばした最短距離の候補を配列に重ね、その最小値をとる、という手順で答えを集める。

-- distance ignoreing
    dI = accumArray min maxBound (2, pred n)
      [ (k,d)
      | (i,s) <- zip [1..n] ss, let d1 = succ $ dF ! i
      , (j,'1') <- zip [i+1..i+m] $ BS.unpack s, let d = d1 + dR ! j
      , k <- [succ i .. pred j]
      ]

距離が巨大すぎるときは、つまりそのような経路がないということなので、それを -1 にする後処理をして完成。

    answer = map post $ elems dI

    post x = if tooBig <= x then -1 else x

結果。

感想

今回は非常にスムーズに考えられた（時間内にできるとは言っていない）