ABC282 A～E をHaskellで

A - Generalized ABC

問題 ABC282A

結果

main = readLn >>= putStrLn . flip take ['A'..]

B - Let's Get a Perfect Score

問題 ABC282B

シグネチャを決める。

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

特定の二人の組が全てを解けるには、全ての問題について、二人のどちらかは解けることが必要。そのような組み合わせを総当たりで試して数える。

結果

import Data.List

abc282b :: Int -> Int -> [String] -> Int
abc282b n m ss = length
  [() | (si,ss1) <- tails ss, sj <- ss1, and $ zipWith canSolve si sj]

canSolve 'x' 'x' = False
canSolve  _   _  = True

$1 \leq M \leq 30$ なので、解けない問題を1とするビット列でそれぞれを表現し、ANDをとってもオール0にならないとき全問正解できない、とすればかっこいいが、$2 \leq N \leq 30$ なのでそこまで速度を稼がなくてもTLEしたりしないので、それを書く時間がもったいない。

import Data.List
import Data.Bits

abc282b :: Int -> Int -> [String] -> Int
abc282b n m ss = length
  [() | b1:bs <- tails (map s2b ss), b2 <- bs, b1 .&. b2 == 0]

s2b :: String -> Int
s2b = foldr step 0
  where
    step 'x' b = shift b 1 .|. 1
    step 'o' b = shift b 1

C - String Delimiter

問題 ABC282C

シグネチャを決める。

abc282c :: Int     -- N
        -> String  -- S
        -> String  -- 答え

ダブルクオートの範囲内と範囲外という二つの状態を持つ有限オートマトンを考えればよい。

範囲外（初期状態）

入力	出力	遷移
"	"	内
,	.	外
c	c	外

範囲内

入力	出力	遷移
"	"	外
c	c	内

結果

abc282c n s = loopO s

-- 範囲外
loopO ('"':s) = '"' : loopI s
loopO (',':s) = '.' : loopO s
loopO ( c :s) =  c  : loopO s
loopO "" = ""

-- 範囲内
loopI ('"':s) = '"' : loopO s
loopI ( c :s) =  c  : loopI s
loopI "" = ""

範囲内かどうかをフラグで持ちまわすスタイルでも書ける。

import Data.List

abc282c :: Int -> String -> String
abc282c n = snd . mapAccumL step False

step b ',' = (b, if b then ',' else '.')
step b '"' = (not b, '"')
step b  c  = (b, c)

D - Make Bipartite 2

問題 ABC282D

シグネチャを決める。

abc282d :: Int          -- N
        -> Int          -- M
        -> [(Int,Int)]  -- ui,vi
        -> Int          -- 答え
abc282d n m uvs = ...

まずいつものように、グラフの辺の情報はノード番号を添え字とする配列に入れておく。

es = accumArray (flip (:)) [] (1,n) [p | (u,v) <- uvs, p <- [(u,v),(v,u)]]

例2の説明に、「与えられるグラフが二部グラフであったり連結であるとは限らない」とある。保証されるのは「単純グラフ」だけである。

全体が二部グラフでないとき、どう辺を足しても二部グラフにはならないので答えは常に0となる。

連結な部分グラフが全て二部グラフをなすとき、異なる部分グラフに属する頂点どうしを接続するとき、それは条件を満たす。
同じ部分グラフに属する頂点どうしについては、両者が異なる色で塗られていれば条件を満たす。

まず、連結な部分を見つけつつ二部グラフになるか調べるために着色することを考える。
適当な頂点から開始して、辺を探索する。隣接する頂点の色を調べ、

• 着色済みの隣接ノードが全て同じ色のとき、その逆の色で塗る。
• 2色に塗り分けられていたら、二部グラフにならず失敗である。

また、未着色の隣接ノードがあれば、それはスタックに積む。

グラフ全体が連結でない可能性もあるので、再帰呼び出しでなく自分でスタックを管理する深さ優先探索を用いて、さらに、スタックの初期状態＝探索予定の開始点を、通常は1頂点のみにする代わりに全ての点を入れる。辺を辿って到達できないノードにはこれを使って訪問するので、最初の頂点を含めて、上の場合分けに次の場合が追加される。これで漏れなく探索される。

• 着色済みの隣接ノードがないとき、そこは別の島なので、新たな色で塗る。

色塗りの深さ優先探索の結果は、着色に成功した場合の、ノード番号から色へのIntMapか、二部グラフでないことが判明した場合の失敗かである。
色は単純に2色ではなく、島ごとに違う2色にすることで、ノードが属する島を見分けることもできるので、色番号/2を島の番号、偶数と奇数で島ごとに2色、という整数で扱う。

import qualified Data.IntMap as IM
import Data.Bits (xor)

type Color = Int
type Node  = Int

mcols :: Maybe (Color, IM.IntMap Color)
mcols = colLoop 0 IM.empty [1..n]

colLoop :: Color                 -- 次に新たな島に到達したときに使う色番号
        -> IM.IntMap Color       -- = Map Node Color ノード番号から色
        -> [Node]                -- 探索スタック
        -> Maybe
             ( Color             -- 最終的に使った色の数
             , IM.IntMap Color)  -- 最終的な着色結果
colLoop cnt im [] = Just (cnt, im) -- スタックが空になったら全て塗れたので成功
colLoop cnt im (i:is)
  | IM.member i im = colLoop cnt im is -- 既に着色されていたらやることなし、次へ
  | otherwise =
-- iの隣接ノードで、未着色なノードのリスト
  let js = [j | j <- es ! i, IM.notMember j im]
-- iの隣接ノードで、着色済みの色リスト
      cs = [c | j <- es ! i, Just c <- [IM.lookup j im]]
  in
    case nub cs of
      []  -> colLoop (cnt+2) (IM.insert i cnt       im) (js++is) -- 新たな島に到達
      [c] -> colLoop  cnt    (IM.insert i (xor c 1) im) (js++is) -- 逆の色で塗る
      _   -> Nothing -- （気持ちは [_,_]）二色と隣接するとき失敗

jsとcsを作るのに2度IntMapを叩くとコスト増になるので、一度だけで済ませるには次のようにする。

import Data.Either
import Data.Maybe

(js,cs) = partitionEithers [maybe (Left j) Right $ IM.lookup j im | j <- es ! i]

塗り分けができたら、全ての頂点間で、辺が存在せず、島の違いも含めて、両端が異なる色になる組み合わせを数え上げる：

import qualified Data.IntSet as IS

(_, cols) = Just mcols

ans = length
  [ ()
  | a <- [1..pred n], let ca = IM.lookup a cols
  , let vs = IS.fromList $ es ! a, b <- [succ a..n], IS.notMember b vs
  , let cb = IM.lookup b cols, ca /= cb]

とやると $O(N^2)$ で間に合わないので、もっと賢く数える必要がある。

• 完全グラフの辺の数は $N(N-1)/2$
• このグラフに既にある辺の数は $M$
• 色 $c$ について、その色で塗られた頂点が $N_C$ 個あるとき、それを繋ぐ $N_C(N_C-1)/2$ 本は除外するべき辺
• それ以外の本数が求める答え

なので、色ごとにその色のノード数を数え、$_{N_C}C_2$ を総数から引く。

Just (cnt, cols) = Just mcols
ca = accumArray (+) 0 (0,cnt) [(c,1) | c <- IM.elems cols]
ans = nC2 n - m - sum (map nC2 $ elems ca)

結果

import qualified Data.IntMap as IM
import Data.Bits (xor)
import Data.Either
import Data.Maybe
import Data.Array

type Color = Int
type Node  = Int

abc282d :: Int -> Int -> [(Int,Int)] -> Int
abc282d n m uvs =
  case mcols of
    Nothing -> 0 -- 二部グラフでない
    Just _  -> nC2 n - m - sum (map nC2 $ elems ca)
  where
    es = accumArray (flip (:)) [] (1,n) [p | (u,v) <- uvs, p <- [(u,v),(v,u)]]
    mcols = colLoop 0 IM.empty [1..n] :: Maybe (Color, IM.IntMap Color)
    Just (cnt, cols) = Just mcols
    ca = accumArray (+) 0 (0, pred cnt) [(c,1) | c <- IM.elems cols]

colLoop :: Color -> IM.IntMap Color -> [Node]-> Maybe (Color, IM.IntMap Color)
colLoop cnt im [] = Just (cnt, im)
colLoop cnt im (i:is)
  | IM.member i im = colLoop cnt im is
  | otherwise =
    case nub cs of
      []  -> colLoop (cnt+2) (IM.insert i cnt       im) (js ++ is)
      [c] -> colLoop  cnt    (IM.insert i (xor c 1) im) (js ++ is)
      _   -> Nothing
  where
    (js,cs) = partitionEithers [maybe (Left j) Right $ IM.lookup j im | j <- es ! i]

nC2 n = div (n * pred n) 2

E - Choose Two and Eat One

問題 ABC282E

シグネチャを決める。

abc282e :: Int    -- N
        -> Int    -- M
        -> [Int]  -- Ai
        -> Int    -- 答え

各ボールを頂点とする完全グラフを考える。
頂点 $i$ から頂点 $j$ への辺には、$(A_i^{A_j} + A_j^{A_i}) \bmod M$ というスコアが割り当てられている。
ボールを食べるとは、このグラフから頂点（とそこに接続する辺）を除くことに相当する。あるボールを除くとき、同時に失われる辺の中の最大スコアを持つ辺を選ぶべきである。
最後に1頂点（これを $R$ とする）が残るだけまでに選んだ辺を考えると、その辺は、$R$ を根とする木をなしているはずである。この木を葉の方から徐々にトリミングしていき、最後に根が残るような順序で、頂点を除いていけるはずである。
ということは、これは元の完全グラフに対して、最大スコアのスパニング木を求める問題に他ならない。
そのような木は、スコア最大の辺から順に、両端が既に連結でなければその辺を使うという貪欲法で、UnionFindを利用して発見できる。（「クラスカル法」という名前が付いている。）

結果

$1 \leq A_i \leq M - 1, 2 \leq M \leq 10^9$ なので、スコアを素朴に (^) と mod で計算するとオーバーフローする。親切にもそれは例2で指摘してくれている。

import Data.List
import qualified Data.IntMap as IM

abc282e :: Int -> Int -> [Int] -> Int
abc282e n m as = loop (pred n) 0 newUF ses
  where
    ses = sortBy (flip compare)
      [(mod (powerish mul 1 a b + powerish mul 1 b a) m, (i, j))
      | (i,a):jbs <- tails (zip [1..] as), (j,b) <- jbs]
    mul x y = mod (x * y) m

loop :: Int               -- 残り使う辺の本数
     -> Int               -- スコア合計
     -> UnionFind         -- 接続状況
     -> [(Int,(Int,Int))] -- 使う辺の候補
     -> Int               -- 総スコア
loop 0 score _ _ = score
loop k score uf ((s,(i,j)):sijs) =
  case uniteUF uf i j of
    Nothing  -> loop k score uf sijs
    Just uf1 -> loop (pred k) (score + s) uf1 sijs

-- @gotoki_no_joe
powerish mul i a b = foldl' mul i [p | (True, p) <- zip bs ps]
  where
    bs = map odd $ takeWhile (0 <) $ iterate (flip div 2) b
    ps = iterate (\x -> mul x x) a

-- @gotoki_no_joe
type UnionFind = IM.IntMap Int

newUF :: UnionFind
newUF = IM.empty

getRoot :: UnionFind -> Int -> (Int, Int)
getRoot uf i =
  case IM.lookup i uf of
    Nothing            -> (i, 1)
    Just k | k < 0     -> (i, - k)
           | otherwise -> getRoot uf k

uniteUF :: UnionFind -> Int -> Int -> Maybe UnionFind
uniteUF uf i j
  | a == b    = Nothing
  | r >= s    = Just $ IM.insert a (negate $ r + s) $ IM.insert b a uf
  | otherwise = Just $ IM.insert b (negate $ r + s) $ IM.insert a b uf
  where
    (a, r) = getRoot uf i
    (b, s) = getRoot uf j

感想

$A_1～A_{N-1}$ と $A_2～A_N$ の上三角行列をどうこうするとか、頂点を食べる順序 $N!$ をDPで構成するとか、勘違いな方向で行き詰ってしまっていた。スパニング木に到達する寸前でも、全頂点を一度ずつ通る経路を探す問題、という微妙に外れたことをしばらく考えていた。