5
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 1 year has passed since last update.

ABC328 A~G をHaskellで

Last updated at Posted at 2023-11-14
  • C : 累積和
  • E, F : Union-Find

宿題:Fを重み付きUnion-Findで書く方法 追記しました。


A - Not Too Hard

問題 ABC328A

シグネチャを決める。

abc328a :: Int   -- N
        -> Int   -- X
        -> [Int] -- Si
        -> Int   -- 答え
abc328a _n x ss = sum $ filter (x >=) ss

問題文の言うとおりにした。

B - 11/11

問題 ABC328B

シグネチャを決める。

abc328b :: Int   -- N
        -> [Int] -- Di
        -> Int   -- 答え

怠惰な解法

カレンダーの全ての日付を生成し、数字に直してみて、ゾロ目なものを数える。

abc328b :: Int -> [Int] -> Int
abc328b n ds = length [() | (i, di) <- zip [1..] ds, d <- [1..di], allSame (show i ++ show d) ]

allSame (x:xs) = all (x ==) xs

考えすぎる解法

可能性のある月は 1~9 と、11の倍数 11~99 のみで、後者の場合、ゾロ目となる数字は 11 で割った商である。
数字 $d$ が決まると、可能性のある日は $d, 11d$ の二つで、$D_i$ を超えない範囲のものを数えればよい。

abc328b :: Int -> [Int] -> Int
abc328b n ds = length [() | (i,di) <- zip [1..] ds, c <- cands i, c <= di]

cands :: Int -> [Int]
cands x
  | x <= 9 = [x, 11 * x]
  | r == 0 = [q, 11 * q]
  | otherwise = []
  where
    (q,r) = divMod x 11

C - Consecutive

問題 ABC328C

シグネチャを決める。$l_i, r_i$は手抜きする。

abc328c :: Int     -- N
        -> Int     -- Q
        -> String  -- S
        -> [[Int]] -- li, ri
        -> [Int]   -- 答え

$S$の先頭から全ての$i$文字目までについて、同じ文字が隣り合う箇所の個数を数えて配列に入れておく。
あとは差を返すだけでよい。

結果

abc328c n _q s lrs = [arr ! r - arr ! l | l:r:_ <- lrs]
  where
    arr = listArray (1,n) $ scanl (+) 0 $ zipWith f s $ tail s
    f c d = if c == d then 1 else 0

D - Take ABC

問題 ABC328D

シグネチャを決める。

abc328d :: String  -- S
        -> String  -- 答え

スタックつきオートマトンのようなものを作ればうまくできる。
'C'に遭遇したときに、現在位置の手前にB,Aがこの順にあったとき、それら全てを消費する。

結果

abc328d :: String -> String
abc328d s = loop "" s
  where
    loop ('B':'A':t) ('C':s) = loop t s
    loop t (c:s) = loop (c:t) s
    loop t "" = reverse t

E - Modulo MST

問題 ABC328E

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

abc328e :: Int     -- N
        -> Int     -- M
        -> Int     -- K
        -> [[Int]] -- ui,vi,wi
        -> Int     -- 答え

木は辺を頂点の個数-1だけ持つので、頂点$N$から始めて、頂点2まで順に辺を1本ずつ、既に連結になっていないものを選ぶ。
最後に、頂点1を含めて連結していれば全域木なのでそのコストを報告する。
最後に、それらの最小値をとる。

結果

Union-Findは今回、uniteUF が新規に連結を行ったとき更新値を、さもなくばNothingを返す形にしている。

abc328e :: Int -> Int -> Int -> [[Int]] -> Int
abc328e n _m k uvws = minimum $ loop newUF 0 n
  where
    g = accumArray (flip (:)) [] (1,n)                       -- いつものグラフ
        [p | u:v:w:_ <- uvws, p <- [(u,(v,w)),(v,(u,w))]]
    loop uf cost 1 = [mod cost k | snd (getRoot uf 1) == n]  -- 頂点1が全部でN個と連結なら総コストを返す
    loop uf cost a =
      [ d
      | (b,c) <- g ! a                      -- 頂点 a と隣接する頂点 b で
      , Just uf1 <- [uniteUF uf a b]        -- 既に到達可能でなかったものについて
      , d <- loop uf1 (cost + c) (pred a)   -- 再帰的にコストを求める
      ]

提出

F - Good Set Query

問題 ABC328F

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

abc328f :: Int     -- N
        -> Int     -- Q
        -> [[Int]] -- ai,bi,di
        -> [Int]   -- 答え

整数列 $(X_1, \dots, X_N)$ それぞれに頂点を対応させたグラフを考える。
問題の設定は、クエリ $(a, b, d)$ を前から順に、$X_a - X_b = d$ が成り立つものを集めると言っている。
このクエリは、$X_a$ と $X_b$ の差を指定していて、それが矛盾していれば採用しない。

とすると、$(a,b)$ を辺とすると、連結でない限り、矛盾が生じることはない。
既に連結されている者同士を次のクエリの辺が接続しようとするとき、矛盾があれば却下せよ、ということ。
なのでまず、初めて接続する辺だけからなる森を作る。
Union-Findで森の木ごとに代表点$r$を決め、$X_r = 0$ と固定し、連結する他の頂点の値を辺から木の探索で決める。

こうして$X_i$の値が全て決定されたら、それに矛盾しない$S_i$の$i$を全て出力する。

immutableなHaskellに翻訳

Union-FindはE問題と同じ。

abc328f :: Int -> Int -> [[Int]] -> [Int]
abc328f n _q abds =
-- Xa - Xb = d であるような i を順に挙げる
    [i | (i, a:b:d:_) <- zip [1..] abds, xm IM.! a - xm IM.! b == d]
  where
-- Union-Findで、スパニング木の森を作り、その辺の情報を取り出す
    (uf, es) = mapAccumL step newUF abds
    step uf abd@(a:b:_) =
      case uniteUF uf a b of
        Nothing  -> (uf , Nothing)
        Just uf1 -> (uf1, Just abd)

-- スパニング木の森なグラフ
    g = accumArray (flip (:)) [] (1,n)
        [p | Just (a:b:d:_) <- es, p <- [(a,(b,-d)),(b,(a,d))]]

-- それぞれの木の代表頂点
    roots = IS.fromList $ map (fst . getRoot uf) [1..n]

-- 代表はX=0とし、gをtraverseして他のXの値を全て求める
    xm = loop IM.empty [] (IS.elems roots)
    loop im [] [] = im
    loop im [] (r:rs) = loop im [(r,0)] rs
    loop im ((i,x):ixs) rs
      | IM.member i im = loop im  ixs  rs
      | otherwise      = loop im1 ixs1 rs
      where
        im1  = IM.insert i x im
        ixs1 = [(b, x + d) | (b,d) <- g ! i] ++ ixs

TLE x 3 で終わった。

mutableな計算に翻訳

さすがにF問題に immutable な Union-Find は舐めプすぎたようなので、同じアルゴリズムを IOVector で再実装したら、 385ms, 217MB でACした。

G - Cut and Reorder

問題 ABC328G

シグネチャを決める。

abc328g :: Int    -- N
        -> Int    -- C
        -> [Int]  -- Ai
        -> [Int]  -- Bi
        -> [Int]  -- 答え

頭に入ってこない問題文

整理してみた。

長さ$N$の数列$A=(A[1],...,A[N]), B=(B[1],...,B[N])$が与えられる。
目的は、$A$に下記の操作を任意回行って$B$と等しくすることである。
それに必要な合計コストの最小値を求めよ。

  • 操作1: 一つの要素の値を任意に変更する。
    整数$i$と整数$k$を任意に選び、$A'[i] = A[i] + k, A'[j] = A[j] \ (j \neq i)$ とした $A'$を新しい$A$とする。
    コストが $|k|$ かかる。
  • 操作2: 数列を任意の断片に分割し、並べ替える。
    $X$個の断片に分割して並べ替えるとは、
    分割の位置 $1 = i[1] < i[2] < \dots < i[X] < i[X+1] = N+1$ と$(1,2,...,X)$の順列$(p[1],\dots,p[X])$を自由に選び、
    $A$の$X$個の断片 $AD[k] = (A[i[k]],A[i[k]+1],...,A(i[k+1]-1))$ $(1\leq k \leq X)$を$AD[p[1]], \dots, AD[P[X]]$ の順に連結したものを新しい$A$とする。
    コストが $(X-1)C$ かかる。

問題の内容は理解できた。
操作1は、切った箇所の個数でコストがかかるので、最初に一度だけ行えば十分である。$A$をいい感じに切って並び替えて、pointwiseに差がなるべく小さくなるようにして、差の合計+切った箇所の個数×$C$を最小にすることが目的である。

そして、$N \leq 22$とか一見ずいぶん小さな問題に見えるところがとても曲者なのだろう。

休むに似たり

$N \leq 22$なので、$A$や$B$の使用済みの値を記憶する整数集合をビット表現で表したりするのだろうと(ズルい)考え方をして、次のようなDPをしてみた。

集めるDPでは表しにくそうなので、諦めて配るDPをする。
$B$を前から順に調べていく。
状態として、$A$の使用済みの要素の番号のビット配列と、直前の$B$を消すのに使った$A$の番号の対について、最小コストを記録しておく。
(巡回セールスマン問題のときに使うようなDP)
次に$B[j]$に(未使用の)$A[i]$を対応させるとき、コストは$|B[j] - A[i]|$と、直前に使った$A$が$i-1$でないときは追加で$C$かかる。
$B$を全て調べた後の、コストの最小値を答えとする。

import Data.Bits
import qualified Data.Map as M

abc328g :: Int -> Int -> [Int] -> [Int] -> Int
abc328g n c as bs = minimum $ M.elems mN
  where
    mN :: M.Map (Int,Int) Int
    mN = foldl' step (M.fromList [((-1,0),0),((-2,0),-c)]) bs
    step m bi = M.fromListWith min
      [ ((j, setBit bmp j), cost + abs (aj - bi) + if succ i == j then 0 else c)
      | ((i, bmp), cost) <- M.assocs m
      , (j, aj) <- zip [0..] as
      , not $ testBit bmp j
      ]

小さな入力については正解するものの、$N=22$な例3についてはどうにもならなかった。

フレンズさんの解説

アライさんの解説
$N$要素のビット配列で、$A$のうち使用済みの分を表す。そのビットの1の個数だけ、$B$は前から順に消費する。
そのような状況の最小コストをDPし、$B$の続きと対応させる$A$の区間を総当たりする。

ということでやってみる。
DP配列の代用としてIntMapを使い、まず0にコスト0を入れた初期値を用意する。
とここで、ステップごとに連続する範囲を割り当てるたびに$C$を足すと、切れ目の数でなく断片の個数になってしまう。これを補填するため、初期値は$-C$にする。
そして、0から順に$2^N-2$までをより添え字の大きい要素に配り、配り終えた後の要素$2^N-1$が答え。
配り終えた要素はその後は不要なので捨ててやると、メモリの節約になるだろうか。

ここまでをコードにする。

import Data.Bits
import qualified Data.IntMap as IM

abc328g :: Int -> Int -> [Int] -> [Int] -> Int
abc328g n c as bs = im IM.! all1
  where
    all1 = 2 ^ n - 1
    im = foldl' step (IM.singleton 0 (negate c)) [0 .. pred all1]
    step :: IM.IntMap Int -- ビットマップ→コスト 配列
         -> Int           -- 今回配る添え字 = Aの使用済み要素ビットマップ
         -> IM.IntMap Int -- 更新された配列
    step im bmp = IM.delete bmp im1 -- 最後に、配り終えた要素を削除(必須ではない)
      where
        im1 = ...                   -- bmpの要素を配り終えた更新済みの配列

コスト配列の注目要素を配るステップ動作stepを考える。
添え字bmpの二進数での1の個数がBの前から使用済み要素数で、
今回はその次からとAのあらゆる(未使用の)箇所とを対応させる。
まずはその部分を取り出す。

        bs1 = drop (popCount bmp) bs -- 今回対応を探すBの部分

配るコストの基本値は、配列の要素に$C$を足した値。

        cost = im IM.! bmp + c -- 現状のコスト(今回のCを含む)

$A$の開始点は、添え字を$1$から$N$まで全て振り、未使用な要素全てを考える。
$A$のその位置から先と、$B$の差を順に求めておく。

              [ ...
              | i <- [0 .. pred n], not $ testBit bmp i
              , let ds = zipWith (\b a -> abs $ b - a) bs1 $ drop i as

$A$のiから始まる断片について、未使用な範囲全ての長さについて、配る先と配るコストを計算する。
範囲を1ずつ広げて、bmpのビットを立てることを累積すると配る先のビットマップが作れる。
costdsの値を順に足し込むと、配るコストが作れる。
終わるのは、bmpのビットが1になったときである。

              , let is = takeWhile (not . testBit bmp) [i .. pred n]  -- i以降でbmpが1になるまでの添え字の列
              , let bmps = scanl setBit bmp is                    -- 配る先のビットマップの列に直す
              , let cs   = scanl (+) cost ds                      -- dsを累積して配るコストの列を作る

このbmpsの要素とcsの要素をそれぞれ組にしてimに戻す。ただし先頭は今注目している位置なので捨てる。

        im1 = foldl' (\im (b,c) -> IM.insertWith min b c im) im
              [ p
              | ...
              , p <- tail $ zip bmps cs
              ]

完成。例1,2は通るので正しそう。
しかしもちろん、$2^{22}$ 要素も抱え込むIntMapは時間も空間も間に合わない。

命令的な計算を導入

同じアルゴリズムを、mutable vectorで実装する。
上のim1のリスト内包表記を極力残す形で書いてみた。

import Data.Bits
import qualified Data.Vector.Unboxed.Mutable as MUV
import Control.Monad

abc328g :: Int -> Int -> [Int] -> [Int] -> IO Int
abc328g n c as bs =
  do
    costs <- MUV.replicate (succ all1) maxBound
    MUV.write costs 0 (negate c)
    forM_ [0..pred all1] (\bmp -> do
      cost <- (c +) <$> MUV.read costs bmp
      forM_ (distrib bmp cost) (\(b1,c1) -> MUV.modify costs (min c1) b1)
      )
    MUV.read costs all1
  where
    all1 = 2 ^ n - 1 :: Int
    distrib bmp cost =
      [ p
      | let bs1 = drop (popCount bmp) bs
      , i <- [0 .. pred n], not $ testBit bmp i
      , let cs   = scanl (+)   cost $ zipWith (\b a -> abs $ b - a) bs1 (drop i as)
      , let bmps = scanl setBit bmp $ takeWhile (not . testBit bmp) [i .. pred n]
      , p <- zip bmp cs
      ]

しかしTLE,MLEする。

空間計算量は、配列の要素数で$O(2^N)$である。
時間計算量は、bmpのループは$2^N$回、distribの中で開始点iを$1$から$N$まで回し、終了点も$N$まで探しに行く可能性があるので、$O(N^2 2^N)$で、これはフェネックさんのいうところの『時間O(N^2 2^N)の想定TLE解法』なのか?

でも、公式解説の日本語の説明はやっぱり目が滑るけれど、実装例(これもC++のなんだか一般的でなさそうな書き方がいっぱいで目が滑るが)のforは3重になっていて、やっていることは同じに見える。けれど、$O(N 2^N)$だと言っている。そしてまたその説明がよくわからん…

配る側でなく、配られる側から考えて、いくつの値が配られるかを考える。
ある状態$S$に集まってくる値は、長さ$1$から長さ$N$まで、何らかの長さ$k$の区間について、$S$で$1$になっている連続区間。
全ての$S$について、ある長さ$k$の1になっている区間を選び出す場合の数は、「連続する$k$個の1」という棒一本と、任意の「0のタイル」「1のタイル」を並べて長さ$N$の数字列を作る方法に等しいので、$2^{N-k}$通り。
ひとつの、長さ$k$の$A$の区間が$B$に対応する位置は$N - k + 1$箇所ある。
よって、配る回数=配られる回数の合計は、$\sum_{k=1}^N 2^{N-k} (N - k + 1) = 2^N (N-1)$ なので、$O(N 2^N)$ なるほど。

計算

$\sum_{k=1}^N 2^{N-k} (N - k + 1) = \sum_{j=0}^{N-1} 2^j (j+1) = T[N]$ とおく。
$T[N] = 2^0 * 1 + 2^1 * 2 + 2^2 * 3 + ... + 2^{N-1} N$
$2T[N] = 2^1 * 1 + 2^2 * 2 + ... + 2^{N-1} (N-1) + 2^N N$
上から下を引くと左辺は $T[N] - 2T[N] = - T[N]$
右辺は $2^0 + 2^1 + 2^2 + ... + 2^{N-1} - 2^N N = 2^N - 1 - 2^N N = 2^N (1 - N) - 1$

結果

distribの内包表記を諦めて、全てを命令的なループで計算する。

abc328g :: Int -> Int -> [Int] -> [Int] -> IO Int
abc328g n c as bs =
  do
    costs <- MUV.replicate (succ all1) maxBound
    MUV.write costs 0 (negate c)
    forM_ [0..pred all1] (\bmp -> do
      let bs1 = drop (popCount bmp) bs
      cost <- (c +) <$> MUV.read costs bmp
      forM_ (filter (not . testBit bmp) [0 .. pred n]) (\i -> do
        foldM_ (\(bmp1, cost1) (i1, ai1, bj1) -> do
          let bmp2 = setBit bmp1 i1
          let cost2 = cost1 + abs (bj1 - ai1)
          MUV.modify costs (min cost2) bmp2
          return (bmp2, cost2)
          ) (bmp, cost) (zip3 (takeWhile (not . testBit bmp) [i .. pred n]) (drop i as) bs1)
        )
      )
    MUV.read costs all1
  where
    all1 = 2 ^ n - 1 :: Int

これなら間に合った。1752ms 44MB
なんかくやしい。
($A$を配列でアクセスするともうすこし速くできる。1152ms 44MB


追記(11/16)

E 問題について

Shuffleは興味深いけれど、$_nC_k$ を計算する、mutable arrayがあるときに有用な、コンパクトで風変わりな方法という以上の意味を見いだせなかった。

たかだか8要素なので、Union-Findでなく、隣接行列をビット表現してN-1乗まですることで全域木か判定してみたが、Union-Findの方が速い感じだった。

F 問題について

重み付きUnion-Find

「重み」というより「高低差」と呼ぶ方がイメージがつきやすいかも。

Union-Findは、それぞれのノードについて、親のノードの番号を表にする。
それに加えて、親のノードとの高低差も、表で記録しておく。

type UnionFind = (IM.IntMap Int, IM.IntMap Int) -- 親頂点、親との高低差

newUF :: UnionFind
newUF = (IM.empty, IM.empty)

分割の代表元を取り出す計算で、親のノードを順に辿るとき、高低差も同時に累積することで、代表元との高低差も取り出せる。

-- getRoot ufht i = (a,h) : iの代表元aと、aからiの高さh
getRoot :: UnionFind -> Int -> (Int, Int)
getRoot (uf, ht) i = loop i 0
  where
    loop j h =
      case IM.lookup j uf of
        Just k -> loop k (h + ht IM.! j)
        _      -> (j, h)

統合を行う計算では、まず、与えられた二つのノードについて、その代表元とそこからの高低差を得る。
両者の代表元が同一なら、高低差の基準が同じなので、そのまま差をとると、与えられた高低差が正しいか判定できる。
両者の代表元が一致しないとき、指定された高低差になるように、代表元どうしの高低差を決めて接続する。

-- uniteUF uf i j d : iとjを連結にし、iからjの高さをdにする 既に連結のときは、高さが一致するか判定
uniteUF :: UnionFind -> Int -> Int -> Int -> (UnionFind, Bool)
uniteUF ufht@(uf, ht) i j d
  | a == b    = (ufht, hbj - hai == d) -- iとjは連結で、iからjの高さはhbj - hai
  | otherwise = ((uf1, ht1), True)     -- aを代表元にし、bを接続する aからbの高さを設定する
  where
    (a,hai) = getRoot ufht i -- hai : aからiの高さ
    (b,hbj) = getRoot ufht j -- hbj : bからjの高さhbj
    uf1 = IM.insert b a uf
    ht1 = IM.insert b (hai + d - hbj) ht

(裏で)経路圧縮を行い親を付け替えるとき、同時に高低差も正しい値を保つように調整する。
これを使うと、この問題の本体部分は極めて簡潔にできる。

abc328f :: Int -> Int -> [[Int]] -> [Int]
abc328f n _q abds = loop newUF 1 abds
  where
    loop _ _ [] = []
    loop uf i ((a:b:d:_):abds) =
      case uniteUF uf b a d of
        (uf1, True ) -> i : loop uf1 (succ i) abds -- 高低差が矛盾しないか、新規の連結のときは番号を出力
        (uf1, False) ->     loop uf1 (succ i) abds -- 高低差に矛盾がある番号は出力しない

G問題について

空間計算量の改善 by hirayuu_AtのMLE版のコードを読み解こうとしたのだが、
$DP1[S][last]$ を設定している 31 行め

				dp[i][j+1]=min(dp[i^(1<<j)][j],mini[i^(1<<j)]+C)+abs(A[b-1]-B[j]);

と、説明

$DP1[S][last]=S$ に含まれる要素までをすでに確定させていて、最後の要素を $last$ で使うときのコストの最小値

から、添え字の意味がどうにも理解できなくて、降参。
(last = j について考えているときに、j+1 について考える、しかもj+1がS = iに含まれるかどうかを無視して何を計算している?)

ここさえ超えれば、

  • 通常のビットDPでは、ビット表現を整数の昇順で計算していけば、注目している状態の部分集合は全て計算済みという性質だけを使うため、それら全体をメモリに保持し、整数のカウントアップの順に計算していく。
  • ここでの手法は、要素が1だけ少ない=ビットが1つだけ落ちた要素だけを参照するので、popCountの数ごとにステップを分けて、不要になったステップの情報は捨てる。

というアイデアでメモリの節約をしているのだろうと推測したのだけれど。

5
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
5
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?