ABC277 A～E をHaskellで

A - ^{-1}

問題 ABC277A

タイトルは「逆関数」と読むのかな。
というか前回のAとかなり被っているような。

シグネチャを決める。

abc277a :: Int    -- N
        -> Int    -- X
        -> [Int]  -- Pi
        -> Int    -- 答え

背番号を（前に）つけてから $X$ を探せばよい。

結果

abc277a n x = fst . head . filter ((x ==).snd) . zip [1..]

(追記)
Data.List.elemIndex :: Eq a => a -> [a] -> Maybe Int を使えば自分で探さずに済む。 （使わないから存在をすっかり忘れていた。）

import Data.List
import Data.Maybe

abc277a :: Int -> Int -> [Int] -> Int
abc277a n x = succ . fromJust . elemIndex x

B - Playing Cards Validation

問題 ABC277B

シグネチャを決める。

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

$N \leq 52$ で大した数ではないので、効率は気にせず計算してしまう。

結果

import Data.List

abc277b :: Int -> [String] -> Bool
abc277b n ss = cond1 && cond2 && cond3
  where
    cond1 = all (isOneOf "HDCS") $ map head ss
    cond2 = all (isOneOf "A23456789TJQK") $ map (!! 1) ss
    cond3 = length (nub ss) == n

isOneOf xs = \x -> elem x xs -- isOneOf = flip elem

C - Ladder Takahashi

問題 ABC277C

シグネチャを決める。

abc277c :: Int          -- N
        -> [(Int,Int)]  -- Ai,Bi
        -> Int          -- 答え

フロアをノード、はしごを辺とするグラフで、1階を含む連結成分で最大のフロアを知りたい。
UnionFindを使いたいが、フロア上限が $10^9$ と大きいので、普通に配列でやるとメモリがあふれる。なので座標圧縮をかける、というのが一つの策なのだろう。

IntMap版のUnionFindで、辺の両端に出現しないノードは生成しないままにしておく版を使えば、直接適用できる。

結果

例3のように、1階が辺に現れていないと、UnionFindに1が現れず、maximum の対象が空になってしまうので、1階は強制的に追加した。

• 検索対象を i <- 1 : IM.keys uf とする
• ダミーのはしごを (0,1) : abs と追加して、1階が必ず現れるようにする

としても対処できる。

import qualified Data.IntMap as IM

abc277c :: Int -> [(Int,Int)] -> Int
abc277c n abs = maximum (1 : [i | i <- IM.keys uf, findUF uf 1 i])
  where
    uf = foldl uniteUF newUF abs

-- UnionFind 実装は略
type UnionFind = IM.IntMap Int
newUF   :: UnionFind
findUF  :: UnionFind -> Int -> Int -> Bool
uniteUF :: UnionFind -> (Int, Int) -> UnionFind

提出版はこちら。

D - Takahashi's Solitaire

問題 ABC277D

シグネチャを決める。

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

それぞれのカードの枚数を配列で数えようとすると、$M \leq 10^9$ と大きいのでメモリが足らなくなる。IntMap でそのまま進んでもよいが、リストだけで解決できる。

カードを書かれている数の順に整列し、数が同じか1だけ大きいものをグループにする。
このとき、0の書かれたカードと$M-1$の書かれたカードがある場合は、両者のグループをひとつにする。
グループのカードに書かれた数の和をそれぞれ求め、その最大値を選ぶ。このグループをなくすように操作するのが最善となる。

結果

0から$M-1$までのカードが全てある場合が、回り込みがある場合と紛らわしく、回り込みとしては処理できないので、先に判定する。

import Data.List

abc277d n m as
  | singleton sums = 0
  | rounded   = sum as - maximum sums1
  | otherwise = sum as - maximum sums
  where
-- 整列
    sas = sort as
-- 回り込みがあるか
    rounded = (head sas == 0) && (last sas == pred m)
-- 続きをグループにし、グループごとの和をとる
    sums = map sum $ foldr step [] sas
-- 回り込みがある場合
    sums1 = (head sums + last sums) : tail (init sums)

-- 次のカードと同じグループに入れられるならそうする。
-- 離れすぎならば新規のグループにする。（次がない、末尾のカードの場合を含む）
step x (ys:yss)
  | head ys - x <= 1 = (x:ys) : yss
step x yss = [x] : yss

-- グループが一つなら、それを取り除くだけ
singleton [_] = True
singleton _  = False

E - Crystal Switches

問題 ABC277E

シグネチャを決める。

abc277e :: Int    -- N
        -> Int    -- M
        -> Int    -- K
        -> [(Int,Int,Int)]  -- ui,vi,ai
        -> [Int]  -- si
        -> Int    -- 答え

辺の長さはどれも1なので、ダイクストラ法を持ち出したりする必要はなく、幅優先探索で着実に距離を数えればたどり着ける。

初期状態のグラフを「モード1」、スイッチを奇数回操作して反転した状態のグラフを「モード0」と呼ぶことにする。両方のモードで個別に、到達済みのノード集合、新規に到達したノード集合を管理する。また、現在までのステップ数をカウントする。

新規に到達したノードがどちらも空なら、進展がなくなったので -1 で終了する。
いずれかの新規到達ノードに N があれば、現在のステップ数が答えである。
どちらでもないとき、新規ノードを到達済みノードに加え、さらに辺を辿って新たなノードに新規到達するが、ここでスイッチが影響する。

新規到達ノードがスイッチのある場所のとき、スイッチを押すことでモードを切り替えることができる。（切り替えなくてもよい。）よって、スイッチの場所の新規到達ノードは、逆側のモードの新規到達ノードでもあると考えられる。

結果

import Data.Array
import qualified Data.IntSet as IS

abc277e :: Int -> Int -> Int -> [(Int,Int,Int)] -> [Int] -> Int
abc277e n m k uvas ss = loop 0 IS.empty IS.empty (IS.singleton 1) IS.empty
  where
-- モード1,0のグラフ
    g1 = accumArray (flip (:)) [] (1,n) [p | (u,v,1) <- uvas, p <- [(u,v),(v,u)]]
    g0 = accumArray (flip (:)) [] (1,n) [p | (u,v,0) <- uvas, p <- [(u,v),(v,u)]]
-- スイッチのあるノード
    sw = IS.fromList ss

-- カウント、モード1,0訪問済み、モード1,0新規到達
    loop cnt b1 b0 n1 n0
      | IS.null n1 && IS.null n0 = -1        -- 失敗
      | IS.member n n1 = cnt                 -- ゴールした
      | IS.member n n0 = cnt
      | otherwise = loop (succ cnt) b1a b0a n1b n0b   -- 次に進む
      where
-- スイッチに乗った逆側を新規に加える
        n1a = n1 ⋃ (n0 ⋂ sw)
        n0a = n0 ⋃ (n1 ⋂ sw)
-- 訪問済みに加える
        b1a = b1 ⋃ n1a
        b0a = b0 ⋃ n0a
-- 新規から辿れる先で未到達の場所が次の新規
        n1b = IS.fromList [k | j <- IS.elems n1a, k <- g1 ! j, k ∉ b1a]
        n0b = IS.fromList [k | j <- IS.elems n0a, k <- g0 ! j, k ∉ b0a]

-- ちょっとした遊び
(⋃) = IS.union
(⋂) = IS.intersection
(∉) = IS.notMember

グラフが多く枝分かれしているとそこでのステップの計算が多くなったりするが、結局、全てのノードに2度到達すれば終わるので、計算量は $O(N \log N)$ 程度で済みそう。
逆に、命令型の更新可能な配列でノード集合を管理すると $O(N^2)$ になりそう。

愚痴

モード1でノード1からノードNと各スイッチまでの距離はダイクストラ法で求められる。
Nまで行ければよし、行けない場合、全てのスイッチとノードNの間の距離を、モード1とモード0の両方についてダイクストラ法で求めて、小さい方を使ってグラフを作り、このグラフの上で改めてダイクストラ法で1からNまでの距離を求める、などとやっているうちに時間切れになってしまった。動かしたところでTLEどまり。考えすぎた。