ABC395をHaskellで

A - Strictly Increasing?

問題 ABC395A

シグネチャを決める。

abc395a :: Int   -- N
        -> [Int] -- Ai
        -> Bool  -- 答え
abc395a _n as = and $ zipWith (<) as (tail as)

リストの全ての隣接に対して何かするときの慣用句。

B - Make Target

問題 ABC395B

シグネチャを決める。

abc395b :: Int   -- N
        -> [String] -- 答え

マス $(i,j)$ は結局どうなるか、を一気に求める関数を立てることもできるはずだが、
考えるのが面倒なので、指示をちょっとモノグサして枠線だけ引く。

結果

import Data.Array.Unboxed

abc395b :: Int -> [String]
abc395b n = [[arr ! (i,j) | j <- [1 .. n]] | i <- [1 .. n]]
  where
    arr :: UArray (Int,Int) Char
    arr = accumArray (flip const) '.' ((1,1),(n,n))
      [ (p, '#')
      | i <- [1, 3 .. div (succ n) 2], let j = succ n - i
      , x <- [i .. j]
      , p <- [(i,x),(x,i),(j,x),(x,j)]
      ]

公式解説の解法2

によると

$\min(i,j,N+1-i.N+1-j)$ が奇数ならば黒、偶数ならば白

とのこと。

abc395b n =
  [ [ if b then '#' else '.'
    | j <- [1 .. n]
    , let b = odd $ minimum [i, j, succ n - i, succ n - j] ]
  | i <- [1 .. n]]

C - Shortest Duplicate Subarray

問題 ABC395C

シグネチャを決める。

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

$A_i$ の値ごとに、それを見た最後の位置 $i$ を記録する。
次にまた同じ値が出現したとき、その間の距離が答えの候補となる。
それらを集めて、最小値を求める。

結果

素直に IntMap を使う版

import qualified Data.IntMap as IM

abc395c :: Int -> [Int] -> Int
abc395c _n as
  | null cands = -1
  | otherwise  = minimum cands
  where
    cands = loop IM.empty $ zip [1 ..] as
    loop _ [] = []
    loop im ((i, a):ias) =
      case IM.lookup a im of
        Nothing -> loop (IM.insert a i im) ias
        Just j  -> succ i - j : loop (IM.insert a i im) ias

$1 \leq A_i \leq 10^6$ と範囲が案外狭かったので配列でやる力まかせ版、というか accumArray で曲芸をやる版

import Data.Array

abc395c :: Int -> [Int] -> Int
abc395c _n as
  | null cands = -1
  | otherwise  = minimum cands
  where
    tooBig = 10^7 :: Int
-- fstは前回の出現位置、sndは答えの候補の最小値
    arr :: Array Int (Int,Int)
    arr = accumArray f (- tooBig, tooBig) (1,10^6) $ zip as [1 ..]
    f :: (Int,Int) -> Int -> (Int,Int)
    f (j, s) i = (i, min s $ succ i - j)
-- 候補
    cands = [s | (_, s) <- elems arr, s < tooBig]

公式解説のやり方

数ごとに出現位置を集めて、隣接するものの距離を数えろ、という、むしろHaskellでやるべきアプローチが筆頭で解説されていた。

import Data.Array

abc395c :: Int -> [Int] -> Int
abc395c _n as
  | null cands = -1
  | otherwise  = minimum cands
  where
-- それぞれの数字の出現位置、降順
    arr = accumArray (flip (:)) [] (1, 10^6) $ zip as [1 ..]
-- 候補
    cands = [succ i - j | is <- elems arr, (i,j) <- zip is $ tail is]

一度貯め込む分遅くなるかと思って技巧を凝らしたのだが、結果としてはこの3つではこれが一番速かった。

その公式解説の末尾にさらっと書かれている

いわゆる尺取り法と呼ばれるテクニックを用いてもこの問題を解くことができます。

は、上二つのアプローチのことだろうか。尻尾が$10^6$本生えた虫…

D - Pigeon Swap

問題 ABC395D

CPUのレジスタリネーミングに通じるところがある内容（？）

命令型配列の書きかえを繰り返す問題なので、シグネチャを決めない。

N 軒の鳩の巣が物理的に並んでいると想定して、位置により背番号 x が付いているとする。
これらには住所があり、番号 1～N の書かれた表札がそれぞれ割り当てられている。
区画整理で住所の番号がたびたび入れ替わる。
問題文でいう「巣 i」とは、表札の番号が i であるような巣のこと、とする。
表札について、二つの配列を維持する：

• $\textit{nest}[x] = m$ ：背番号 x の巣には表札は m が掲げられている
• $\textit{addr}[m] = x$ ：住所 m の表札は背番号 x の巣に掛かっている

はとについて、配列を維持する：

• $\textit{pige}[i] = x$ ：鳩 i は背番号 x の巣にいる

クエリには次のように対応する

• 種類1：鳩 a の居所を、表札 b の巣にする： $\textit{pige}[i] \leftarrow \textit{addr}[b]$
• 種類2：住所 a と b の表札を入れ替える：
  aとbの位置を $x = \textit{addr}[a], y = \textit{addr}[b]$ として
  $\textit{nest}[x] \leftrightarrow \textit{nest}[y], \textit{addr}[a] \leftrightarrow \textit{addr}[b]$
• 種類3：鳩 a のいる巣の表札を出力する $\textbf{out}(\textit{nest}[\textit{pige}[a]])$

結果

swap があるので Vector を使った。
AC, 152ms

E - Flip Edge

問題 ABC395E

シグネチャを決める。手抜きする。

abc395e :: [Int]   -- N,M,X
        -> [[Int]] -- ui, vi
        -> Int     -- 答え

表世界として、定義通りのグラフを考える。辺のコストは全て 1 とする。
裏世界として、辺の向きを逆にしたグラフを考える。辺のコストは 1。
実は両者は繋がっており、対応する頂点間をコスト X で行き来できる。

というグラフを構築し、頂点1から頂点Nまたは裏のNまでの距離をダイクストラ法で求めればよい。

結果

表の頂点番号 $1$～$N$ に対して、裏には $N+1$～$2N$ を割り当てる。
関数 gf は、頂点番号に対して、出る辺の先と重みの対のリストを返すことで
グラフの構造を関数に伝えるコールバック関数である。
通常は配列 g に全て仕舞っておいて (g !) を渡すだけだが、
このようにその場で作る使い方もできる。

import Data.Array

abc395e :: [Int] -> [[Int]] -> Int
abc395e [n,_m,x] uvs = runST $
  do
    ds <- dijkstra (1, n + n) gf 1
    dn <- readArray ds n        -- Nまでの距離
    dnn <- readArray ds (n + n) -- 2Nまでの距離
    return $ min dn dnn
  where
    g = accumArray (flip (:)) [] (1, n + n) $
        [(u, v) | u:v:_ <- uvs] ++        -- 正方向
        [(n + v, n + u) | u:v:_ <- uvs]   -- 逆方向
    gf i
      | i <= n    = (n + i, x) : [(j, 1) | j <- g ! i] -- 表から裏へ
      | otherwise = (i - n, x) : [(j, 1) | j <- g ! i] -- 裏から表へ

全文

F - Smooth Occlusion

問題 ABC395F

シグネチャを決める。

abc395f :: Int     -- N
        -> Int     -- X
        -> [[Int]] -- Ui, Di
        -> Int     -- 答え

全くわからない。

二分探索

アライさんからのヒント1は公式解説のやり方。
高さ $0 \leq H \leq \max(U_i + D_i)$ で条件を満たす最大値を二分探索で探す。

元の高さが $U_i+D_i$ 高さを $H$ にするので切り詰める高さは $C_i = U_i+D_i-H$
上の歯の可能な範囲を探す。
最も短くなる場合は、可能な限り上の歯で、足りない分は下の歯で、とするときで、上の歯は $\max(0, U_i - C_i) = \max(0, H - D_i)$
最も長くなる場合は、可能な限り下の歯で、足りない分は上の歯で、とするときで、
下の歯で足りない分とは $\max(0, C_i - D_i) = \max(0, U_i - H)$
これを上の歯から切るので $U_i - \max(0, U_i - H)$

一つ手前の歯の高さが $\textit{lb}$ 以上 $\textit{ub}$ 未満のとき、差が $X$ 以下という条件から、
$[\textit{lb} - X, \textit{ub} + X] \cap [\max(0, H - D_i), U_i - \max(0, U_i - H)]$ が最終的に許される範囲。
これが途中で空になるとき、不可能。全ての歯について取り得る値があるなら成功。

abc395f :: Int -> Int -> [[Int]] -> Int
abc395f n x uds = sum (concat uds) - n * h
  where
    u1 = head $ head uds
    maxh = minimum $ map sum uds
    (_,h) = binarySearch prop (succ maxh) 0
    prop h = loop 0 u1 uds
      where
        loop _l _u [] = True
        loop lb ub ([ui,di]:uds) =
            lb1 <= ub0 && lb0 <= ub1 &&
            loop (max lb0 lb1) (min ub0 ub1) uds
          where
            lb0 = max 0 $ lb - x
            ub0 = ub + x
            lb1 = max 0 $ h - di
            ub1 = ui - max 0 (ui - h)

O(N)の解法

フェネックさんからのヒント2はより高速な解法 by MMNMMのやり方。
しかし

～ となることが示せます。

がさっぱりわからない。
意味はさっぱり分からないまま、コードだけ移植してみた。

import Data.List

abc395f :: Int -> Int -> [[Int]] -> Int
abc395f n x uds = sum (concat uds) - n * minimum hs
  where
    [u1, d1] = head uds
    (_, hs) = mapAccumL step (u1, d1) uds
    step (u, d) [ui, di] = ((u1, d1), u1 + d1)
      where
        u1 = min ui $ u + x
        d1 = min di $ d + x

G - Minimum Steiner Tree 2

問題 ABC395G

シグネチャを決める。

abc395g :: Int     -- N
        -> Int     -- K
        -> [[Int]] -- Cij
        -> Int     -- Q
        -> [[Int]] -- s_i, t_i
        -> [Int]   -- 答え

ABC364Gと同じく（タイトルにもある）シュタイナー木の問題。
基本的には前と同じアルゴリズムを使う。
今回、$\textit{Cost}[X][v]$ の $X$ を振る範囲が $\bigcup_{K < s \leq N} \{1, \dots, K, s\}$ と、番号の飛んでいるところがある点がやっかい。

公式解説のPythonコードでは、1～Kに関する計算を行13～31のループで行い、
これを使って、個々の値 $s$ に関する計算を行33～55の、ほぼコードクローンなループで行っている。
違いは行17に対応する行41で、分割ステップの片方だけを前段の計算結果から持ってくるようにしている。
このため、分割ステップを $(X,\emptyset)$ ～ $(\emptyset,X)$ の $2^{|X|}$ やる必要があるし、これで正しいことも納得しにくい。

集合のビット表現を拡張して、 $\bigcup_s \{1, \dots, K, s\}$ とその全ての部分集合に対して
隙間なく背番号を割り当てることができれば、このような小細工なく実装できる。
これはそう難しくない。
$K$以下の要素だけからなる集合$X$にビット表現を割り当てる関数を$F(X)$とする。
$K$超の要素を含まないか、たかだか一つを含む集合$Y$に対して拡張ビット表現を割り当てる関数$G(Y)$は、
$Y \subseteq \{1, \dots, K\}$ のとき $G(Y) = F(Y)$
$Y \ni s > K$ のとき $G(Y) = F(Y \cap \{1,\dots,K\}) + 2^{|K|} \times (s - K)$
もっと端的に言うと、整数のビット0～ビット$K-1$に要素1～$K$を割り当て、
ビット$K$以上を整数として解釈して0のとき$s$はなし、1以上の値 $i$ のとき $K+i$ が $Y$ に含まれているとする。

また、分配ステップを前回はダイクストラ法で実装したが、今回は完全グラフなためにキューに投入される要素が多くなり、
キューの先頭が何になるのかを毎回計算し直す $O(N^2)$ のやり方の方がよい、と公式解説にあるので従う。

実装

import Data.Array.Unboxed

abc395g :: Int -> Int -> [[Int]] -> Int -> [[Int]] -> [Int]
abc395g n k css _q sts = ...
  where
    tooBig = div maxBound 8 :: Int
    reg x = min tooBig x
-- 完全グラフ
    c = listArray ((1,1),(n,n)) $ concat css :: UArray (Int,Int) Int

ノード1～Kにビット0～$K-1$ を対応させる。

import Data.Bits

type BitSet = Int

K以下全て

    core :: BitSet
    core = pred $ bit k

要素数を数える

    bsSize :: BitSet -> Int
    bsSize bs
      | bs <= core = popCount bs
      | otherwise  = succ $ popCount (bs .&. core)

最小の番号のノード番号（1始まりに戻る）

    theBit bs
      | bs <= core = popCount $ bs .^. pred bs
      | otherwise = k + (bs .>>. k)

分割を全て列挙する

    powerIntSet :: BitSet -> [(BitSet,BitSet)]
    powerIntSet is
      | is <= core = ps
      | otherwise  = concat [[(a .|. s, b), (a, b .|. s)] | (a,b) <- ps]
      where
        is1 = is .&. core
        ps = takeWhile (uncurry (>)) $ loop is1
        loop x = (x, is1 .^. x) : loop (is1 .&. pred x)
        s = is - is1

K以下の要素だけからなるビット集合に、K超の要素を入れた背番号を求める insert 関数

    bsInsert :: Int -> BitSet -> BitSet
    bsInsert i bs
      | k < i = bs + ((i - k) .<<. k)
      | True  = error "never"

Cost配列
添え字の上限は ${1,\dots,K,N}$ の背番号

    ub = bsInsert n core
    cost :: Array Int (UArray Int Int)
    cost = listArray (0, ub) $ map build [0 .. ub]

各要素を求める関数 build
$X$ の要素数が0,1,2以上、により動作が切り替わる。
（※：ABC364Gと同一）

    build :: BitSet -> UArray Int Int
    build is =
      case bsSize is of
        0 -> listArray (1, n) $ repeat 0
        1 -> sending $ accumArray (flip const) tooBig (1, n) [(theBit is, 0)]
        _ -> sending $ splitting is

分割ステップは、powerIntSet で得られる分割のうち、先頭以外について自己参照した結果を足し合わせる。
（※：ABC364Gと同一）

    splitting bs = accumArray min tooBig (1, n)
      [ (v, c1 + c2)
      | (is1, is2) <- tail $ powerIntSet bs
      , ((v, c1), c2) <- zip (assocs $ cost ! is1) (elems $ cost ! is2)]

分配ステップは、現時点で未調査かつ距離最小の点からの辺で残りを更新する、を最大$N$回繰り返す。

import qualified Data.IntSet as IS

    one2n = IS.fromList [1 .. n]
    sending :: UArray Int Int -> UArray Int Int
    sending im = loop one2n im
      where
        loop is im
          | v == 0    = im
          | otherwise = loop is1 im1
          where
            (cost, v) = minimum $ (tooBig, 0) : [(c, v) | (v,c) <- assocs im, IS.member v is]
            is1 = IS.delete v is
            im1 = accum min im [(u, cost + c ! (v,u)) | u <- IS.elems is1]

$\forall K<s\leq N; \textit{Cost}[{1,\dots,K,s}][v]$ が求まれば、$s = s_i$ としたときの $v = t_i$ の値が答えである。

abc395g :: Int -> Int -> [[Int]] -> Int -> [[Int]] -> [Int]
abc395g n k css _q sts = [(cost ! bsInsert s core) ! t | s:t:_ <- sts]
  where
    ...

提出：2061ms