ABC391をHaskellで

A - Lucky Direction

問題 ABC391A

シグネチャを決める。

abc391a :: String -- D
        -> String -- 答え

文字ごとに処理する方法

DにあるNとS、EとWを互いに交換すればよい。

abc391a = map oppo

oppo 'N' = 'S'
oppo 'S' = 'N'
oppo 'E' = 'W'
oppo 'W' = 'E'

方角として扱う方法

方角を表す文字列のリスト ["N","NE","E","SE","S","SW","W","NW"] を考える。
ある方角の反対方向は、リストの4つ後（ループしているとみなす）にある。

import Data.Maybe

abc391a :: String -> String
abc391a d = fromJust $ lookup d oppo

dirs = ["N","NE","E","SE","S","SW","W","NW"]
oppo = zip dirs $ drop 4 $ cycle dirs

B - Seek Grid

問題 ABC391B

シグネチャを決める。

abc391b :: Int    -- N
        -> Int    -- M
        -> [String] -- Sij
        -> [String] -- Tij
        -> (Int,Int)-- 答え

総当たりで探すしかない。
座標の足し算を分かりやすくするために、Tは0始まりで扱う。

結果

import Data.Array.Unboxed

abc391b :: Int -> Int -> [String] -> [String] -> (Int,Int)
abc391b n m ss ts = head
  [ (i, j)
  | i <- [1 .. n - m1], j <- [1 .. n - m1]
  , and [s ! (i + p, j + q) == tpq | ((p,q),tpq) <- assocs t] ]
  where
    m1 = pred m
    s, t :: UArray (Int,Int) Bool
    s = listArray ((1,1),(n,n))   $ map ('#' ==) $ concat ss
    t = listArray ((0,0),(m1,m1)) $ map ('#' ==) $ concat ts

C - Pigeonhole Query

問題 ABC391C

シグネチャを決める。

abc391c :: Int     -- N
        -> Int     -- Q
        -> [[Int]] -- query_i
        -> [Int]   -- 答え

複数の視点で状況を追跡する。

• それぞれの鳩がどの巣にいるか
• それぞれの巣には鳩が何羽いるか
• 複数の鳩がいる巣の個数はいくつか

結果

IntMap を immutable な配列の代用品にしても間に合う。1145ms, 186MB
ST の MVector を使って、上のコードとほぼ同じ形で255ms, 108MB
IOUArray を使って、main で全部やってしまう形で152ms, 40MB

それにしても immutable な Vector の modify 関数は
本当に in-place でやれるならやる機構を持っているのだろうか？

D - Gravity

問題 ABC391D

シグネチャを決める。

abc391d :: Int     -- N
        -> Int     -- W
        -> [[Int]] -- Xi, Yi
        -> Int     -- Q
        -> [[Int]] -- Ti, Ai
        -> [Bool]  -- 答え

全てのブロックの高さをtickごとに落としていくなんてことはもちろんできない。
地面の方がtickごとにせり上がっていき、最下行が全部揃ったら次のtickで消える。
ブロックが消えるのは、最下行が全員揃ったときのみ。
これは、列ごとに見て、前から k 番目という順番が同じブロックどうしで消えること、その順番の前から順に消えること、がわかる。
また、ブロックが十分スカスカなら、同じ順番のブロックのうち初期位置がもっとも高いブロックが地面に落ちたときが消えるときだとわかる。
さらに、消えずにブロックが積もり、床のせり上がりによりブロックが押し戻される状況を追跡しなくても、初期位置で最も高い位置のブロックが、その高さぶんのtickの経過で必ず地面に到達することもわかる。
（みっしり詰まってせり上がる列と、スカスカなせいで消えない列があったとして、それでも、スカスカな列のブロックが下まで落ちるには高さ分のtickしかかからない。
消えるブロックの組は下からの順番が同じものなので、せり上がる列の対応する高さのブロックを、スカスカな列のブロックが追い越そうとしても、そこより下には同じ数だけブロックがある上にスカスカで空間が空いているので、追い越すことはできない。伝わるだろうか…）

結局、列ごとに、ブロックを高さの低い順に並べ、全体について、下からの順番が同じものどうしを、高さとブロック番号の情報を集める。
横一列W個揃っている行までは消えて、ブロックの足りない行以上は消えない。
一行のブロックは同時に消え、その消えるタイミングは最も高いブロックの高さと等しい。

これでそれぞれのブロックが消える時刻がわかるので、クエリに答えることができる。

結果

import Data.Array.IArray
import Data.List

abc391d :: Int -> Int -> [[Int]] -> Int -> [[Int]] -> [Bool]
abc391d n w xys q tas = [t < tarr ! a | t:a:_ <- tas]
  where
-- 列ごとにブロックを、高さをキー背番号を値とするIntMapに登録
    columns :: Array Int [(Int, Int)]
    columns = amap sort $ accumArray (flip (:)) [] (1,w) [(x, (y,i)) | (i, x:y:_) <- zip [1 ..] xys]
-- ブロックの最も少ない列の行までが消える
    h = minimum $ map length $ elems columns
-- 下からの高さごとに、最も高いブロックの高さの時刻に行全員が同時に消える、その時刻をブロックに対して記録する
    tarr :: Array Int Int
    tarr = accumArray (flip const) maxBound (1, n)
      [ (j, t)
      | yis <- take h $ transpose $ elems columns
      , let t = maximum $ map fst yis
      , j <- map snd yis
      ]

E - Hierarchical Majority Vote

問題 ABC391E

シグネチャを決める。

abc391e :: Int    -- N
        -> String -- Ai
        -> Int    -- 答え

0から1なのか、1から0なのかにはこだわらず、「結果をひっくり返す」ことだけに注目する。
再帰的な構造になっていて、一番下は、一人の票をひっくり返すのに必要なコストは1だが、
それより上の段について、ある一人をひっくり返すには、それぞれの結果をひっくり返すのに必要なコストがまるまるかかる。それらを $c,d,e$ とする。
一番下の段は $c = d = e = 1$ である。以降は、下の段のコストを持ち上げて考える。

3人の投票	結果	結果を変えるのに必要な操作人数
000	0	誰か二人を変えるので $c + d + e - \max(c,d,e)$
001	0	0のうち一人を変えるので $\min(c,d)$
010	0	同上で$\min(c,e)$
100	0	同上で$\min(e,f)$
110	1	1のうち一人を変えるので $\min(c,d)$
101	1	同上で$\min(c,e)$
011	1	同上で$\min(e,f)$
111	1	誰か二人を変えるので $c + d + e - \max(c,d,e)$

投票結果とその票をひっくり返すコストの対のリストを作り、
3つごとに区切っては上の計算を行い、
結果をつなぎ合わせたリストに繰り返し同じ事を行うと、
最後に長さ1のリストに答えが残る。

結果

import Data.List.Split

abc391e :: Int -> String -> Int
abc391e n as = loop [(a, 1) | a <- as]

loop [(_,c)] = c
loop acs = loop $ map step $ chunksOf 3 acs

step ((a,x):(b,y):(c,z):_) =
  case [a,b,c] of
    "000" -> ('0', x + y + z - maximum [x,y,z])
    "001" -> ('0', min x y)
    "010" -> ('0', min z x)
    "011" -> ('1', min y z)
    "100" -> ('0', min y z)
    "101" -> ('1', min z x)
    "110" -> ('1', min x y)
    "111" -> ('1', x + y + z - maximum [x,y,z])

公式解説は、かなり抽象化した始点から説明しているのか、なんかピンとこないけど、同じ事をやっているような気はする。

F - K-th Largest Triplet

問題 ABC391F

シグネチャを決める。

abc391f :: Int   -- N
        -> Int   -- K
        -> [Int] -- Ai
        -> [Int] -- Bi
        -> [Int] -- Ci
        -> Int   -- 答え

大きい方から K 個を数え上げる方法

A,B,Cを整列させて考える。
$(i,j,k)$ の座標に $f(i,j,k) = A_iB_j + B_jC_k + C_kA_i$ を配置すると、添え字が大きくなるほど値も大きくなるから、$(n,n,n)$ の値が最大で、面で隣接するマスの値は、これを使った後でだけ考慮すればよい。

なので、これらのマスのどれを使用したか、というフラグ配列を準備して、優先度を $f(i,j,k)$ 値を $(i,j,k)$ とする優先度付きキューに入れて、大きい順に値を取り出し、取り出すたびに最大3つの隣接マスのうち未登録なものをキューに追加する、を K 回繰り返せばよいはず…$N \le 2 \times 10^5$ に対して $N^3$ 要素の配列は張れない。

アプローチ間違ってるのかとフレンズさんのヒントを見ると

アライグマ「F問題はプライオリティキューを使うと解けるのだ！

合ってそう。でもわからなくて解説を見た。やはり方針はあってた。

ポイントは、面で接するマスが次の候補になっていくとき、$(n,n,n)$に近い位置から使われていき、決して、使用済みの位置がえぐれることはない、ということ。$(i,j,k)$ よりも先に $(i,j-1,k)$ が使われたりはしない。
これを裏返すと、「登録済みかもしれないマスは、登録済みかもしれないが、使用済みにまでなっていることはない」ということ。つまり「登録が重複しない仕組みがあれば、重複して登録してしまえばよい」「なので、登録済みかどうかを記録しておく必要はない」

Data.Heap は重複した登録を防ぐ仕組みはない。
代用として Data.Set.deleteFindMax を使うと、重複した登録はちょうど無視できる！

import Data.List
import Data.Array.Unboxed
import qualified Data.Set as S

abc391f :: Int -> Int -> [Int] -> [Int] -> [Int] -> Int
abc391f n k as bs cs = loop k q0
  where
    aA, bA, cA :: UArray Int Int
    aA = listArray (1,n) $ sort as
    bA = listArray (1,n) $ sort bs
    cA = listArray (1,n) $ sort cs
    f (i, j, k) = aA ! i * (bj + ck) + bj * ck
      where
        bj = bA ! j
        ck = cA ! k
    q0 = S.singleton (f (n,n,n), (n,n,n))
    loop 1   q = fst $ S.findMax q
    loop cnt q = loop (pred cnt) q3
      where
        ((_, (i,j,k)), q1) = S.deleteFindMax q
        ijks = [(pred i,j,k) | 1 < i] ++ [(i,pred j,k) | 1 < j] ++ [(i,j,pred k) | 1 < k]
        q3 = S.union q1 $ S.fromList [(f ijk, ijk) | ijk <- ijks]

改善版

公式解説1の補足 by kyopro_friendsの工夫を使うと、登録が重複する心配をしなくもよくなる、つまり Data.Heap のままでもできるという。
やってみたところ、優先度付きキューを Data.Set のままにしたら1565ms、Data.Heapにしたら2415msでむしろ遅くなってしまった。

二分探索をする版

優先度付きキューで K 個数えるのでなければ二分探索か、とやってみたが自力ではTLEを解消できなかった。
公式 解説2 by toamではできていて、二分探索の述語を手抜きで $O(N^3)$ かかるような書き方をしていたらダメで、$O(K)$ にできるとあった。

数えている途中で K 個になったらループを脱出する、という手続き的な計算になっているのでどうしたものかと思ったが、何とかなった。ポイントは sumIsGE 関数。

import Data.List
import Data.Array.Unboxed

abc391f :: Int -> Int -> [Int] -> [Int] -> [Int] -> Int
abc391f n k as bs cs = snd $ binarySearch prop (succ $ g a1 (bA ! 1) c1) 5
  where
    aA = listArray (1,n) $ sortBy (flip compare) as :: UArray Int Int
    bA = listArray (1,n) $ sortBy (flip compare) bs :: UArray Int Int
    cA = listArray (1,n) $ sortBy (flip compare) cs :: UArray Int Int
    a1 = aA ! 1
    c1 = cA ! 1
    g a b c = a * (b + c) + b * c
-- x以上の値がk個以上
    prop x = sumIsGE k
        [ 1
        | b  <- takeWhile (\b -> g a1 b c1 >= x) $ elems bA
        , c  <- takeWhile (\c -> g a1 b c  >= x) $ elems cA
        , _a <- takeWhile (\a -> g a  b c  >= x) $ elems aA ]

-- 非負の数リストxsの総和がk以上かどうかを判定する
sumIsGE k xs
  | k <= 0    = True
  | null xs   = False
  | otherwise = sumIsGE (k - head xs) (tail xs)

-- @gotoki_no_joe
binarySearch :: (Int -> Bool) -> Int -> Int -> (Int, Int)
binarySearch prop unsat sat = loop unsat sat
 where
   loop a b
     | ende   = (a, b)
     | prop m = loop a m
     | True   = loop m b
     where
       ende = a == m || b == m
       m = div (a + b) 2

G - Many LCS

問題 ABC391G

シグネチャを決める。

abc391g :: Int    -- N
        -> Int    -- M
        -> String -- S
        -> [Int]  -- 答え

最初の答案

具体的な2つの文字列に対する編集距離またはLCSの長さを求めるのに、
AtCoderの解説では二次元配列を広げるDPと説明される。
これを、注目している行のひとつ上の行の内容だけ持っておけばできる、はわかっている。

あとは、注目している行に対応する T の文字の全ての可能性に対して、
上の行の状況の全ての組み合わせとその場合の数に対して、次の状況と場合の数を計算し直す、という
外側のDPをかぶせればできる、と理屈では思いついたが、状況の種類が多すぎて爆発して終わりではないか、と納得できず着手できなかった。

フレンズさんいわく

アライグマ「G問題は、LCSを求めるDPの配列を状態に持ったDPをするのだ！

合ってた。ならやってみよう。

行の状況をキー、その場合の数を値とする Map を外側のDPの状態とする。
キーのメモリ表現をなるべくコンパクトにするため、Data.Vector.Unboxed を使う。

import Data.Array.Unboxed
import qualified Data.Vector.Unboxed as UV
import qualified Data.Map as M
import Data.List

abc391g :: Int -> Int -> String -> [Int]
abc391g n m s = elems ans
  where
-- Sに現れる文字リストと倍数の対
    charsmags = zip ('*' : nub s) (starmag : repeat 1)
-- Sに現れない文字の種類数
    starmag = 27 - length charsmags
-- 初期状態
    st0 = M.singleton (UV.fromList $ replicate (succ n) 0) 1 :: M.Map (UV.Vector Int) Int
-- 最終状態
    stM = foldl' step st0 [1 .. m]
-- 状態と場合の数カウント、という大状態を、Tをもう1文字追加する遷移関数
    step m _ = M.fromListWith add
      [ (next st c, mul cnt fac)
      | (st, cnt) <- M.assocs m
      , (c, fac) <- charsmags ]
-- 状態とTの文字について、stの次の行を作る
    next st c = UV.fromList st1
      where
        st0 = UV.toList st
        st1 = 0 : [if c == si then succ c7 else max c4 c8 | (si,c4,c7,c8) <- zip4 s st1 st0 (tail st0)]
-- 最終状態のLCSの値をキーにして、答えをまとめる
    ans = accumArray add 0 (0, n) [(st UV.! n, cnt) | (st, cnt) <- M.assocs stM] :: UArray Int Int

p :: Int
p = 998244353
reg :: Int -> Int
reg x = mod x p
add :: Int -> Int -> Int
add x y = reg $ x + y
mul :: Int -> Int -> Int
mul x y = reg $ x * y

できた！825ms, 117MB と思ったら、C++勢は10ms少々で終わっている。速すぎない？

もっと軽く

ちゃんと公式解説を読むと、LCSの長さは最大でも$\min(N,M)$なので、$N$を越えることはない。
作り方から、状況の配列の内容は、先頭が0、そこから1増えるか増えないかの単調増加にしかならないと。
つまりこれを長さ $N+1$ の Vector Int64 で格納するのはぜいたくすぎで、増えるか増えないかの1ビットを $N$ 個の1ワードで十分だと。その表現なら配列の添え字にできるので、上の Map は配列に置き換えられる。

結果：321ms, 12MB
上のコードよりだいぶ速くなったし、メモリ消費が圧倒的に減ったので、こんなとこでしょう。