ABC309 A～F をHaskellで

A - Nine

問題 ABC309A

シグネチャを決める。

abc309a :: Int   -- A
        -> Int   -- B
        -> Bool  -- 答え

結果

真面目にやるなら、$A+1=B$ と $A \not\in \{3,6,9\}$ を調べる。

abc309a a b = succ a == b && mod a 3 /= 0

正解の組み合わせは6通りなので、それを全て列挙しても大したことない。

abc309a a b = elem (a,b) [(1,2),(2,3),(4,5),(5,6),(7,8),(8,9)]

B - Rotate

問題 ABC309B

シグネチャを決める。

abc309b :: Int       -- N
        -> [String]  -- Aij
        -> [String]  -- 答え

取り出す位置は、3とおりの場合に分けて考えることができる。

• 辺の要素は、１つ「手前」の位置から取り出す。
• ただし、先頭は「手前」がないので、「前」の辺の末尾と考え、ここでは扱わない。
• それ以外の内側は、その位置から取り出す。

結果

import Data.Array

abc309b :: Int -> [String] -> [String]
abc309b n ass = [[arr ! f i j | j <- [1..n]] | i <- [1..n]]
  where
    arr = listArray ((1,1),(n,n)) $ concat ass
    f i j
      | i == 1, j /= 1 = (i, pred j) -- 上辺
      | j == n, i /= 1 = (pred i, j) -- 右辺
      | i == n, j /= n = (i, succ j) -- 下辺
      | j == 1, i /= n = (succ i, j) -- 左辺
      | otherwise      = (i,j)       -- 内側

C - Medicine

問題 ABC309C

シグネチャを決める。

abc309c :: Int          -- N
        -> Int          -- K
        -> [(Int,Int)]  -- ai, bi
        -> Int          -- 答え

積分としての累積で関数を復元する。
初日は $\sum b_i$ 錠飲む必要がある。
$a_i+1$日目には、前日より $b_i$ 錠は減る。
数の範囲が大きいので、配列でなくIntMapで扱う。

import qualified Data.IntMap as IM

abc309c :: Int -> Int -> [(Int,Int)] -> Int
abc309c _n k abs = succ ans
  where
    sumb = sum $ map (!! 1) abs
    im = IM.fromListWith (+) abs
    (ans, _) = head $ dropWhile ((k <) . snd) $ scanl step (0, sumb) $ IM.assocs im
    step (_, acc) (a, b) = (a, acc - b)

D - Add One Edge

問題 ABC309D

シグネチャを決める。

abc309d :: Int -> Int   -- N1,N2
        -> Int          -- M
        -> [(Int,Int)]  -- ai, bi
        -> Int          -- 答え

番号1の頂点と連結な$N_1$頂点のグラフと、番号$N_1+N_2$の頂点と連結な$N_2$頂点のグラフがあり、この２つは繋がっていない。
どこか一カ所だけ繋いで、番号1と番号$N_1+N_2$が一番遠くなるようにしたときの距離はいくつか、と聞かれている。

元の状態で1から最も遠い頂点までの距離$D_1$を求め、
同様に$N_1+N_2$から最も遠い頂点までの距離$D_2$を求め、
$D_1+D_2+1$が答え。

辺の長さは全て1なので、ダイクストラ先生を煩わせずともBFSで数えられる。

結果

import Data.Array
import qualified Data.IntSet as IS

abc309d :: Int -> Int -> Int -> [(Int,Int)] -> Int
abc309d n1 n2 m abs = 1 + maxdist g 1 + maxdist g n
  where
    n = n1 + n2
    g = accumArray (flip (:)) [] (1, n)
        [p | (a,b) <- abs, p <- [(a,b),(b,a)]]

maxdist g v0 = loop 0 IS.empty [v0] []
  where
    loop d _ [] [] = max 0 (pred d)
    loop d visited [] news = loop (succ d) visited news []
    loop d visited (v:vs) news
      | IS.member v visited = loop d visited vs news
      | otherwise = loop d (IS.insert v visited) vs (g ! v ++ news)

loopのそれぞれの引数の意味を下に示す。

newsの更新では、初到達の頂点vから何も考えずに g ! v を追加しているので、逆に、newsが空になるとは、前回のvsで初到達が１つもなかったことになる。なので1を引いた値を結果とする。
ただし、出発点が辺を持たないという特異な状況で結果が-1になってしまうので、そうならないように最後に補正をかけている。

    loop :: Int    -- 開始点からの現在の距離d
         -> IS.IntSet -- 訪問済みの頂点集合
         -> [Int]  -- 初到達なら距離dとなる、調べる頂点リストvs
         -> [Int]  -- vsに隣接する頂点リスト、次の周回で距離d+1か調べる
         -> Int    -- 結果 v0 から最も遠い頂点の距離

追記：Data.SequenceによるFIFOでBFS

上の loop は２つのリストでぎったんばっこんしてBFSしているが、FIFOを使えばもっと素直にできる。

import qualified Data.Sequence as Q

    maxdist v0 = last $ loop IS.empty (Q.singleton (v0, 0))
    
    loop _ Q.Empty = []
    loop visited ((v,e) Q.:<| q)
      | IS.member v visited = loop visited q
      | otherwise = (e :) $
                    loop (IS.insert v visited) $
                    (q Q.><) $
                    Q.fromList [(w, succ e) | w <- g ! v]

E - Family and Insurance

問題 ABC309E

シグネチャを決める。

abc309e :: Int          -- N
        -> Int          -- M
        -> [Int]        -- Pi
        -> [(Int,Int)]  -- xi, yi
        -> Int          -- 答え

自分の掛けた保険は、効力最大のもの１つだけ考えればよい。
親に掛かっている保険は、効力を1減らしたものが自分にも掛かっていると解釈できる。
効力0までが補償対象者である。

結果

import Data.Array

abc309e :: Int -> Int -> [Int] -> [(Int,Int)] -> Int
abc309e n m ps xys = length $ filter (0 <=) $ elems assu
  where
    ys = elems $ accumArray max (-1) (1,n) xys
    assu = listArray (1,n) $                  -- 3. が自分の保険
           zipWith max ys $                   -- 2. 自分の保険の大きい方
           (-1) : [pred $ assu ! p | p <- ps] -- 1. 親の保険-1と

F - Box in Box

問題 ABC309F

シグネチャを決める。$h_i, w_i, d_i$ は手抜きする。

abc309f :: Int      -- N
        -> [[Int]]  -- hi, wi, di
        -> Bool     -- 答え

似たようなシチュエーションで、最大いくつの箱を重ねられるか、と聞かれる問題はLIS 最長上昇部分列問題だが、これはそっちではない。
一組でいいので、すっぽりはまる箱の対が見つかればよい。

まず、全ての箱の向きを $h_i \leq w_i \leq d_i$ となるように統一して、この向きについてのみ考えればよい。（そうでない向きで収まるなら、この向きでも収まるはずだから。）なので全てソートする。これを$x_i,y_i,z_i$ と呼ぶことにする。

次に、全体を、$x_i$ の大きいものから順に考え、既に調べ終わった箱たちの中に、今注目している箱が収まるか、という流れで進めていく。すると、$x_i$ については等しいかより大きなものだけを調査済みなので、「より大きくて調査済み」の情報の中だけを調べればよい。見つかればその時点で終了し、見つからなかった場合、この箱の $x_i, y_i, z_i$ について登録を追加して次に進む。

調査済みの箱について取っておくべき情報は、それぞれの $y$ における最大の $z$ の値。大は小を兼ねる。
次に調べる$y_i, z_i$ について、$y_i < y$ かつ$z_i < z$ なものが存在すればよい。
そしてこの「それぞれの$y$における最大の $z$」は、「大は小を兼ねる」の観点で単調減少をなす。（あとあまり関係ないが、$z \geq y$ の範囲しか点がない。）

結果

import Data.List
import qualified Data.IntMap as IM

abc309f :: Int -> [[Int]] -> Bool
abc309f _n = loop im0 im0 0 . sortBy (flip compare) . map sort

im0 = IM.fromAscList [(0,tooBig),(tooBig,0)] -- 番兵

tooBig = 10^9 + 1

loop :: IM.IntMap Int  -- xがより大きいものだけのy,zの情報
     -> IM.IntMap Int  -- これまで全てのy,zの情報
     -> Int            -- 前回のx
     -> [[Int]]        -- [[x,y,z]]
     -> Bool           -- 答え
loop _ _ _ [] = False
loop im1 im2 x0 ((x:y:z:_):xyzs) = z < z1 || loop im3 im4 x xyzs
  where
-- xが前回より減少したなら、im1はim2で上書きする
    im3 = if x0 == x then im1 else im2
-- yがより大きいもののzの最大値
    Just (y1,z1) = IM.lookupGT y im3
-- im2の単調性を維持しながらy,zを追加する
    Just (y2,z2) = IM.lookupGE y im2
    im4 = if z <= z2 then im2 else (regularize $ IM.insert y z im2)
-- yがより小さくてzも以下な点を全て削除する
    regularize im
      | z3 <= z   = regularize $ IM.delete y3 im
      | otherwise = im
      where
        Just (y3,z3) = IM.lookupLT y im

公式解説では「座標圧縮してセグメント木で」云々とやっているが、話をややこしくしているだけな気がする…

G - Ban Permutation

問題 ABC309G

公式解説「包除原理を使います。」（り…りろんはしってる）
ユーザ解説「箱根駅伝DPで解釈します。」（えっ何それ）
ユーザ解説「この解法は「合計（数え上げ）」の問題を「期待値」の問題に読み替えるという典型の一種です。」（えっ？えっ？）