ABC302 A～G をHaskellで

• A 切り上げ除算
• B 総当たりの計算
• C グラフの深さ優先探索
• D IntSet.lookupLE （命令型言語では二分探索）
• E ？強いて言えば、命令型配列？
• F 二部グラフの距離（深さ優先探索）
• G ？

A - Attack

問題 ABC302A

シグネチャを決める。

abc302a :: Int -- A
        -> Int -- B
        -> Int -- 答え

つまり切り上げ除算。

結果

abc302a = divrup

divrup x y = negate $ div (negate x) y

B - Find snuke

問題 ABC302B

シグネチャを決める。

abc302b :: Int         -- H
        -> Int         -- W
        -> [String]    -- Si
        -> [(Int,Int)] -- 答え

2次元配列に文字盤を構成しておく。
全ての開始位置から、全ての8方向に向けて、実際に5文字読みだしてみて、すぬけが見つかったらその位置を返す。

文字盤を踏み外さないかのチェックは、最後の5文字めの座標が範囲内にあるかどうかだけ確認すればよい。

結果

import Data.Array

abc302b :: Int -> Int -> [String] -> [(Int,Int)]
abc302b h w ss = ans
  where
    bnds = ((1,1), (h,w))
    ca = listArray bnds $ concat ss
    ans = head
      [ xy15
      | xy1 <- range bnds                         -- 全てのマスから始めて
      , ca ! xy1 == 's'                           -- ただし's'と書いてあるマスだけ
      , d <- dirs                                 -- 全ての方向に
      , let xy15 = take 5 $ iterate (add d) xy1   -- 5文字分の座標列
      , inRange bnds (last xy15)                  -- 末尾まではみ出していない
      , "snuke" == map (ca !) xy15                -- すぬけを探す
      ]

dirs = [(x,y) | x <- [-1..1], y <- [-1..1], x /= 0 || y /= 0]

add (a,b) (c,d) = (a+c,b+d)

C - Almost Equal

問題 ABC302C

シグネチャを決める。

abc302c :: Int       -- N
        -> Int       -- M
        -> [String]  -- Si
        -> Bool      -- 答え

$S_i$と$S_j$の編集距離が1のとき、$i$と$j$を結ぶ辺があるグラフを考える。
このグラフにハミルトン路があるかを探す。
$N \leq 8$なので深さ優先探索で総当たりすればよい。

結果

深さ優先探索において、訪れていない頂点集合をビット集合で表現した。

import Data.Array
import Data.List
import Data.Bits

abc302c :: Int -> Int -> [String] -> Bool
abc302c n m ss = any (dfs b0) [0..n1]
  where
    n1 = pred n
    sis = zip [0..] ss                   -- Siに添え字iを付けて、ただし0始まり
    g = accumArray (flip (:)) [] (0,n1)              -- グラフ
        [ p
        | (i,si):sjs <- tails sis, (j,sj) <- sjs
        , 1 == length (filter id (zipWith (/=) si sj))
        , p <- [(i,j),(j,i)]
        ]
    b0 = bit n - 1 :: Int        -- この型注釈をつけないとData.Bitsに怒られる
    dfs b v                      -- 未訪問ノードのビット集合b, 今踏むノードv
      | b1 == 0   = True
      | otherwise = any (dfs b1) [u | u <- g ! v, testBit b1 u]
      where
        b1 = clearBit b v

D - Impartial Gift

問題 ABC302D

シグネチャを決める。

abc302d :: Int    -- N
        -> Int    -- M
        -> Int    -- D
        -> [Int]  -- Ai
        -> [Int]  -- Bi
        -> Int    -- 答え

$B_i$をどれか選んだとする。これと組み合わせて、条件を満たす範囲で最も価値を高められる$A_i$の選択は、$B_i + D$以下で最大のものとなる。ただし$B_i - D$以上でないとき却下される。
IntSet.lookup(LT|LE|GT|GE)の出番。

結果

候補が一つもなかったときに-1を返す対応は、候補リストに-1を加えておけばよい。

maximum $ -1 : cands

これが「最小値を求めよ。存在しないときは-1を返せ」という問題の場合は、空かどうかを真面目に判断する必要がある。さもないと全ての答えが-1になってしまう。

if null cands then -1 else minimum cands

import qualified Data.IntSet as IS

abc302d :: Int -> Int -> Int -> [Int] -> [Int] -> Int
abc302d n m d as bs = maximum $ -1 : cands
  where
    aS = IS.fromList as
    cands =
      [ b + a
      | b <- bs
      , Just a <- [IS.lookupLE (b + d) aS]
      , b - d <= a
      ]

E - Isolation

問題 ABC302E

辺の削除ができるUnion-Findとか必要か？と泡を喰ってしまったが、よく考えるとそんなにややこしくなかった。

全ての頂点について、辺で結ばれている頂点のリストを持つ、いつものような配列を用いると、計算量が大変なことになるので、頂点の IntSet をもつ配列を考える。
また、クエリに対応してこれを書き変えていくため、mutable な IOArray にする。
辺の配列とともに、「現在、孤立している頂点の個数」も追跡する。初期値は$N$である。

クエリ 1 u v では、$u$ と $v$ に互いを加える。これらが変更前は孤立していたら、孤立頂点の個数をそれだけ減らす。

クエリ 2 v では、$v$ の隣接頂点全てに対して、その隣接頂点から $v$ を取り除く。これで孤立頂点になったものの個数だけ、カウントを増やす。また、$v$ も孤立頂点になるが、元々孤立頂点な場合もあるので注意。

どちらも $O(\log N)$ なので、全体で $O(Q \log N)$ で済む。

結果

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

import Data.Array.IO
import qualified Data.IntSet as IS

main = do
  [n,q] <- bsGetLnInts
  arr <- newArray (1,n) IS.empty :: IO (IOArray Int IS.IntSet)
  foldM (\cnt _ -> do
    qi <- bsGetLnInts
    cnt1 <- (cnt +) <$> case qi of
      1:u:v:_ -> case1 arr u v
      2:v:_   -> case2 arr v
    print cnt1
    return cnt1
    ) n [1..q]

bsGetLnInts :: IO [Int]
bsGetLnInts = unfoldr (BS.readInt . BS.dropWhile isSpace) <$> BS.getLine

case1 :: IOArray Int IS.IntSet -> Int -> Int -> IO Int
case1 arr u v = do
  arru <- readArray arr u
  writeArray arr u $ IS.insert v arru
  arrv <- readArray arr v
  writeArray arr v $ IS.insert u arrv
  return $ 0 - (if IS.null arru then 1 else 0) - (if IS.null arrv then 1 else 0)

case2 :: IOArray Int IS.IntSet -> Int -> IO Int
case2 arr v = do
  arrv <- readArray arr v
  writeArray arr v IS.empty
  delta <- sum <$> forM (IS.elems arrv) (\u -> do
    arru <- readArray arr u
    let arru1 = IS.delete v arru
    writeArray arr u arru1
    return $ if IS.null arru1 then 1 else 0
    )
  return $ delta + if IS.null arrv then 0 else 1

クエリの間で持ちまわす状態がすべてmutableなデータ構造に収まっていれば、繰り返しは replicateM q で簡潔に書けるが、今回のようにimmutableな変数もある場合、

foldM (\状態 _ -> ループ) (初期状態) [1..q]

の形で書いている。もう少し何か気の利いた書き方はないだろうかと思いつつ。

F - Merge Set

問題 ABC302F

シグネチャを決める。$A_i$ は使わないので捨てる。

abc302f :: Int      -- N
        -> Int      -- M
        -> [[Int]]  -- Sij
        -> Int      -- 答え

1から$M$を頂点として、$a, b \in S_i$ のとき $a$ と $b$ を辺で繋ぎ、1から$M$の距離を数えると、いくつの集合を経由するかわかるので、1を引けば融合の回数になる。
…とやると、総当たりな $A_i (A_i - 1) / 2$という辺の本数は爆発を招く恐れがある。

追加で、集合$S_1 ～ S_N$に対応する頂点$M+1 ～ M+N$を置き、$a \in S_i$ のとき $a$ と $M+i$ を繋ぐことにすれば、辺の本数は $A_i$ で済む。
このとき、1から$M$の距離は経由した集合の個数の倍になる。

結果

import Data.Array
import qualified Data.IntMap as IM

abc302f :: Int -> Int -> [[Int]] -> Int
abc302f n m ass = loop IM.empty 0 [1] []
  where
    g = accumArray (flip (:)) [] (1,m+n)
        [ p
        | (i, as) <- zip [succ m..] ass
        , a <- as
        , p <- [(i, a), (a, i)]
        ]
    loop im dist [] ms
      | IM.member m im = pred $ flip div 2 $ im IM.! m
      | null ms        = -1
      | otherwise      = loop im (succ dist) ms []
    loop im dist (n:ns) ms
      | IM.member n im = loop im  dist ns ms
      | otherwise      = loop im1 dist ns ms1
      where
        im1 = IM.insertWith (flip const) n dist im
        ms1 = g ! n ++ ms

深さ優先探索はloopで行っている。

loop :: IM.IntMap Int -- im 頂点の1からの距離
     -> Int           -- dist 現在の距離
     -> [Int]         -- ns 距離distでマークするべき頂点リスト
     -> [Int]         -- ms nsの頂点の隣接頂点リスト、次の周回で探索する
     -> Int           -- 1からMまでの距離/2-1(つまり問題の答え)、または-1

G - Sort from 1 to 4

問題 ABC302G

シグネチャを決める。

abc302g :: Int    -- N
        -> [Int]  -- Ai
        -> Int    -- 答え

$A_i$は1から4しかない。
それぞれの数字の個数を数えることで、整列済みの様子を再現することができ、それと初期状態を比較することで、「本来Aのあるべき位置にBがある個数」16とおりを数え上げることができる。

この表について、

• 対角要素4つについては、正しい位置にあるので、何もしなくてよい。
• 「Aのあるべき位置にあるB」と「Bのあるべき位置にあるA」は互いに交換することで、正しい位置に戻すことかできる。これを全ての組み合わせについて行い、表を更新する。交換の回数も記録する。
• 「Aの位置にあるB」「Bの位置にあるC」「Cの位置にあるA」は、2回の交換により3つの数を正しい位置に戻せる。これの全ての組み合わせは、1から4のうち3つを取りだした順列 $_4P_3$通りある。これも上と同様に表を更新し、交換回数を数える。
• 「Aの位置のB」「Bの位置のC」「Cの位置のD」「Dの位置のA」は3回の交換により4つの数を正しい位置に戻せる。上と同様にする。
  これを上から順に可能なだけやれば、全てが正しい位置に戻るはずである。

結果

表の大きさがたかだか16要素なので、immutableな配列の更新で済ませる。

表の内容を状態に持つStateモナドで、操作によって更新される表の持ちまわしを水面下に隠し、操作の結果としての交換回数の集計が関心事としてコードに表現できた。

import Control.Monad
import Data.List
import Data.Array
import Control.Monad.State

type Mat = Array (Int,Int) Int

abc302g :: Int -> [Int] -> Int
abc302g n as = evalState action mat0
  where
-- 1から4がそれぞれいくつあるか
    cnts = accumArray (+) 0 (1,4) [(a,1) | a <- as]
-- 整列したらどんな姿になるか
    goal = [i | (i,c) <- assocs cnts, _ <- [1..c]]
-- Xのあるべき位置にYがある個数の表
    mat0 = accumArray (+) 0 ((1,1),(4,4)) [(ga, 1) | ga <- zip goal as]
    action = do
      c1 <- sum <$> mapM swap [[a,b] | a <- [1 .. 3], b <- [succ a .. 4]]
      c2 <- sum <$> mapM swap [[a,b,c] | a <- [1..4], b <- [1..4], b /= a, c <- [1..4], c /= a, c /= b]
      c3 <- sum <$> mapM swap (permutations [1..4])
      return $ c1 + c2 + c3

-- [a,b,c]に対して、(a,b)→(b,c)→(c,a) と可能な限り交換する
-- つまりこれらの最小値を全てから引く。
-- (引いた値)×(座標の列の個数-1)を返す。
swap :: [Int] -> State Mat Int
swap is = do
  let ps = zip is (tail is ++ [head is])
  mat <- get
  let x = minimum $ map (mat !) ps
  when (x > 0) $ set $ accum (-) mat [(p, x) | p <- ps]
  return $ (x *) $ pred $ length is

これで、数の上限が4より大きくなったら計算量は爆発する、んだろうな。