A - Conflict
シグネチャを決める。
abc409a :: Int -- N
-> String -- T
-> String -- A
-> Bool -- 答え
同じ位置に o があれば True にする。
結果
要は zip を使えばいい。
abc409a _n t a = any (('o','o') ==) $ zip t a -- 素直な表現
abc409a _n t a = elem ('o','o') $ zip t a -- lintがこうしろと横槍を入れる
abc409a _n t a = or $ zipWith p t a -- タプルを作らずに済ませる
where
p 'o' 'o' = True
p _ _ = False
B - Citation
シグネチャを決める。
abc409b :: Int -- N
-> [Int] -- Ai
-> Int -- 答え
$A$ に $x$ 以上の要素が重複を含めて $x$ 回以上現れる。
という条件を満たす最大の整数 $x$ を見つける。
$A_i$ を降順にソートしたものを $B_i$ とする。
すると、添字 $k$ の値以上の要素は $k$ 回（以上）現れている。
つまり、$i \leq B_i$ を満たす最大の $i$ が答え。
結果
import Data.List
abc409b :: Int -> [Int] -> Int
abc409b n = fst . head . dropWhile (uncurry (>)) . zip [n, pred n ..] . sort
C - Equilateral Triangle
シグネチャを決める。
abc409c :: Int -- N
-> Int -- L
-> [Int] -- di
-> Int -- 答え
大前提として $L$ は3の倍数である必要がある。
$d_i$ を累積し $\bmod L$ した値で円周上の絶対位置 $r_i$ を作り、位置ごとに点の個数を数える。
$0 \leq a < L/3$ に対して、$L/3 \leq b < 2L/3$, $2L/3 \leq c < L$ が正三角形を作るので、その組み合わせの個数を数える。
結果
import Data.Array.Unboxed
import Data.List.Split
abc409c :: Int -> Int -> [Int] -> Int
abc409c n l ds
| r /= 0 = 0
| otherwise = sum $ map product $ transpose $ chunksOf l3 $ elems cnt
where
(l3, r) = divMod l 3
cnt = accumArray (+) 0 (0, pred l) [(a, 1) | a <- scanl add 0 ds] :: UArray Int Int
add x y = let z = x + y in if z >= l then z - l else z -- modなしで済ませる、ケチ臭い高速化
せっかく配列に入れたのだから chunksOf など使わなくてもいい気もする。
zipWith (*) (map (cnt !) [0 .. pred l3]) $
zipWith (*) (map (cnt !) [l3 ..])
(map (cnt !) [l3 + l3 ..])
D - String Rotation
シグネチャを決める。テストケース一つ分の計算をする。
abc409d :: Int -- N
-> String -- S
-> String -- 答え
なるべく早い位置から動かす方が辞書順で影響が大きいので、前の方から可能性を考える。
$S$ の文字を順に $C_1, C_2, \dots$ とし、
$C_i > C_{i+1}$ なら、$i = \ell$ とした結果は辞書順で小さくなる。
$C_i \leq C_{i+1}$ なら、より後方で実施するべき。
選んだ文字 $C_i$ をどこに挿入するべきか。
これも、辞書順での影響が大きくなるように、できるだけ前の方にするべき。
$C_i < C_k$ なら、$r + 1 = k$ とするべき。
$C_i \geq C_k$ なら、もっと後ろにするべき。
結局、$C_i > C_{i+1}$ な最初の文字 $C_i$ を抜いて、
それ以降の $C_i < C_k$ な最初の文字の前、そのような $C_k$ が見つからなかった場合は末尾に入れる。
都合のいい $C_i$ が見つからなかったときは、$\ell = r = 1$ とすることで元の文字列そのままが最小になる。
結果
abc409d :: Int -> String -> String
abc409d _n s = loop1 s
where
loop1 (a:b:cs) | a > b = b : loop2 a cs
loop1 (a:bs) = a : loop1 bs
loop1 "" = ""
loop2 c as@(a:_) | c < a = c : as
loop2 c (a:as) = a : loop2 c as
loop2 c "" = [c]
E - Pair Annihilation
シグネチャを決める。横着する。
abc409e :: Int -- N
-> [Int] -- x_i
-> [[Int]] -- u_i, v_i, w_i
-> Int -- 答え
ABC408F では一直線に行っていた平滑化を、木の上で行うちょっとした変種。
一般のグラフだと気が狂うがここでは木なので、葉から順に寄せて行けばできる。
注目している頂点に対して、その子孫から平滑化したお釣りの電荷が集まってくるので、
この頂点では自身の電荷に加えてそれらを合わせた分を親に向けて平滑化する。
これを流すときに、親との辺の抵抗値を掛けただけの電力が消費される。
なので、dfsでは、消費電力合計と、しわ寄せた電荷の合計の両方を集計する。
結果
import Data.Array
abc409e :: Int -> [Int] -> [[Int]] -> Int
abc409e n xs uvws = ans
where
x = listArray (1,n) xs
g = accumArray (flip (:)) [] (1,n) $
[(u,(v,w)) | u:v:w:_ <-uvws] ++
[(v,(u,w)) | u:v:w:_ <-uvws]
(_0, ans) = dfs 0 1 0
dfs :: Int -- 親
-> Int -- ノード ここを電荷0にする
-> Int -- 親からノードへの辺の抵抗値w
-> ( Int -- 戻ってきた電荷量
, Int) -- 累積した答え
dfs p v wv = (da, sum zs + wv * abs da)
where
(ds, zs) = unzip [dfs v c w | (c, w) <- g ! v, c /= p]
da = sum ds - x ! v
F - Connecting Points
シグネチャを決める。横着する。
abc409f :: Int -- N
-> Int -- Q
-> [[Int]] -- x_i, y_i
-> [[Int]] -- query_i
-> [String] -- 答え
考える
頂点の連結状況を管理するにはUnion-Findを使えばいけるだろう。
次に連結する2点とは、距離が短くまだ連結でないものということなので、優先度付きキューで距離を管理するのだろう。
頂点が追加されるので、既存の頂点の座標も全て覚えておく必要がある。
ということで、以下の状態を考える。
import qualified Data.Heap as PQ
type State =
( Int -- k : 存在する点の個数
, [[Int]] -- xys : 点の座標Xi,Yiの逆順のリスト
, PQ.Heap (Int,Int,Int) -- pq : 距離、点番号ふたつの優先度付きキュー
, UnionFind -- uf : 点番号から連結か判定するUnionFind
)
クエリには次のように対応する。
クエリ1 [1,a,b] に対して、(a,b) を k+1 番めの点として追加する。
また既存の点全て xys との距離を測り、pq に投入する。
クエリ2 [2] に対して、
- キュー先頭のエントリから取り出し、連結なものを全て捨てる
- 連結でないエントリ
(d,i,j)が見つかったら、答えはd - キュー先頭のエントリから取り出し、距離が
dである組を全て Union-Find で連結にする
クエリ3 [3,u,v] に対して、Union-Find に問いあわせた結果が答え
初期の点 N 個も、クエリ1と同じロジックで登録すればよい。
mhd :: [Int] -> [Int] -> Int
mhd (x:y:_) (z:w:_) = abs (x - z) + abs (y - w)
abc409f _n _q xys qus = catMaybes ans
where
st0 = (0, [], PQ.empty, newUF)
st1 = foldl mode1 st0 xys -- XiYiはクエリ1として処理
(_st2, ans) = mapAccumL step st1 qus
step st (1:xy) = (mode1 st xy, Nothing) -- クエリ1は出力なし
step st (2:_) = mode2 st -- クエリ2
step st (3:u:v:_) = (st, Just $ mode3 st u v) -- クエリ3は状態変化なし
mode1 :: State -> [Int] -> State
mode1 (k, xys, pq, uf) xy = (k1, xy:xys, PQ.union pq pq1, uf)
where
k1 = succ k
pq1 = PQ.fromList [(mhd xy zw, k1, j) | (j, zw) <- zip [k, pred k ..] xys]
mode2 :: State -> (State, Maybe String)
mode2 (k, xys, pq0, uf0) = loop1 pq0
where
loop1 pq =
case PQ.uncons pq of -- 距離は短いけど連結な対を捨てる
Nothing -> ((k, xys, pq, uf0), Just "-1")
Just ((d, i, j), pq1) | findUF uf0 i j -> loop1 pq1
| otherwise -> loop2 d pq uf0
loop2 d pq uf =
case PQ.uncons pq of -- 距離 d な対を全て連結する
Just ((e, i, j), pq1) | d == e -> loop2 d pq1 (uniteUF uf i j)
_ -> ((k, xys, pq, uf), Just $ show d)
mode3 :: State -> Int -> Int -> String
mode3 (_,_,_,uf) u v = if findUF uf u v then "Yes" else "No"
immutableなUnion-Findは省略。
提出結果：ACx28, TLEx17
mode2は PQ.dropWhile や PQ.span を使うと自前でループせず書けるが、これをやるとTLEがドンと増えてしまった。後者はともかく前者は納得いかない。
mode2 (k, xys, pq0, uf0)
| PQ.null pq1 = ((k, xys, pq1, uf0), Just "-1")
| otherwise = ((k, xys, pq3, uf1), Just $ show d)
where
pq1 = PQ.dropWhile (\(_,i,j) -> findUF uf0 i j) pq0 -- 距離は短いけど連結な対を捨てる
(d, _, _) = PQ.minimum pq1
(pq2, pq3) = PQ.span (\(e,_,_) -> d == e) pq1 -- 距離 d な対を全て連結する
uf1 = foldl' (flip ($)) uf0 [\uf -> uniteUF uf i j | (_,i,j) <- PQ.toUnsortedList pq2]
純粋計算を捨てて
これ以上の性能を求めるには、immutableなUnion-FindとData.Heapを捨てて、 命令型配列の速度に頼らざるを得ない。
ヒープとしては、AtCoder では Data.Vector.Algorithms.Heapが使える。
なぜか、ベクタに格納する値に対して Ord 型クラスを要求する代わりに、比較関数を引数として毎回渡す形になっている。
type Comparison e = e -> e -> Ordering
compare を渡せばいいように思うが、ひっくり返す flip compare 必要がある。
というのは、これは命令型配列のヒープソートを作ることを目標にしたヒープ機構なので、アルゴリズムの教科書に書いてあるヒープソートのやり方のとおり、配列の前半をヒープ、後半をソート済みのデータにして、ヒープの先頭に浮かび上がった最大値をヒープ末尾の値と交換し、ヒープを1マス縮小することを繰り返し、最終的に昇順ソートされた配列を作るようになっているからである（と想像される。ドキュメントには一言も説明はない。）
具体的な設計に入る。
点を追加したり、ヒープの要素を増減させたりと、配列に実際に入っている要素数が変動するので「末尾に追加する」操作が欲しい。
（C++ の .push_back() メソッドがいつもうらやましい。）
ベクタ自体は可変なデータ構造でも、末尾位置のインデックスを引数で別に取り回すのはスマートでないので、それをバンドルしたデータ構造を作っておく。
import qualified Data.Vector.Unboxed.Mutable as MUV
import Data.IORef
type IOVPlus a = (IORef Int, MUV.IOVector a)
newPlus :: MUV.Unbox a => Int -> IO (IOVPlus a)
newPlus n = do
p <- newIORef 0
v <- MUV.new n
return (p, v)
pushBack :: MUV.Unbox a => IOVPlus a -> a -> IO ()
pushBack (rk, v) x = do
k <- readIORef rk
MUV.write v k x
writeIORef rk (succ k)
setK :: IOVPlus a -> Int -> IO ()
setK (rk, _) = writeIORef rk
getK :: IOVPlus a -> IO Int
getK (rk, _) = readIORef rk
IOVPlus を流用して、データ個数を管理してくれるヒープデータ構造として使えるようにするAPIを追加する、というか Data.Vector.Algorithms.Heap の薄いラッパを被せる。
doHeapify :: (Ord a, MUV.Unbox a) => IOVPlus a -> IO ()
doHeapify (rk, v) = do
k <- readIORef rk
heapify (flip compare) v 0 k
pushPQ :: (Ord a, MUV.Unbox a) => IOVPlus a -> a -> IO ()
pushPQ (rk, v) x = do
k <- readIORef rk
heapInsert (flip compare) v 0 k x
writeIORef rk (succ k)
nullPQ :: IOVPlus a -> IO Bool
nullPQ pq = (0 ==) <$> getK pq
peekPQ :: MUV.Unbox a => IOVPlus a -> IO a
peekPQ (_,v) = MUV.read v 0
popPQ :: (Ord a, MUV.Unbox a) => IOVPlus a -> IO ()
popPQ (rk, v) = do
k <- pred <$> readIORef rk
pop (flip compare) v 0 k
writeIORef rk k
return ()
上の純粋実装では $X_i, Y_i$ を登録するのにクエリ1を流用していたが、これをすると $O(n \log n)$ かかる。
最初のデータは配列に一気に突っ込んで、heapify で一度にヒープ化すると $O(n)$ でできるのでそのための doHeapify である。
またここでは $N$ 個の点間の距離 $N(N-1)/2$ 個を投入するので実際には $O(N^2 \log N)$ から $O(N^2)$ に下がる。
（$N \leq 1500$ だけど。）
あとのアルゴリズムは純粋実装と同じにして、AC, 734ms, 114MB という結果になった。
G - Accumulation of Wealth
公式解説によると convolution を使って計算するらしいが、そこまでに至る議論がまったくギリシャ語なので、読めるようになるまで大事にとっておく。