ABC259 A～E をHaskellで

A - Growth Record

問題 ABC259A

シグネチャを決める。

abc259a :: [Int]   -- N,M,X,T,D
        ->  Int    -- 答え

時計を逆回しすると、

歳	身長	備考
$N$	$T$
$\vdots$	$\vdots$
$X$	$T$	ここまで変化なし
$X-1$	$T-D$	以降、毎年$D$ずつ縮む
$\vdots$	$\vdots$
$0$	$T-XD$

となる。つまり、$X \leq M$ ならば $T$、$X > M$ のときはその差 $\times D$ を引けばよい。
または、差 $X - M$ 掛ける $D$ を $T$ から引くが、差が負のときは0に打ち切る、と考えることもできる。

結果

abc259a [n,m,x,t,d]
  | x <= m = t
  | otherwise = t - d * (x - m)

abc259a [n,m,x,t,d] = t - d * max 0 (x - m)

B - Counterclockwise Rotation

問題 ABC259B

シグネチャを決める。

abc259b :: [Int]             -- A,B,D
        -> (Double,Double)   -- 答え a', b'

普通に回転行列

\left (
\begin{array}{rr}
\cos \theta & - \sin \theta \\
\sin \theta & \cos \theta
\end{array}
\right )

を用いればよい。

結果

abc259b abd = (c * a - s * b, s * a + c * b)
  where
    [a,b,d] = map fromIntegral abd
    theta = d * pi / 180
    c = cos theta
    s = sin theta

C - XX to XXX

問題 ABC259C

シグネチャを決める。

abc259c :: String -> String   -- S,T
        -> Bool               -- 答え

理屈としては、同じ文字が続く区間を group で切り出し、対応する区間同士でその文字は等しく、区間の長さはSで1のときはTでも1でなければならず、Sで1より大きいときは、Tでもそれ以上であればよい。

import Data.List

-- WA
abc259c s t = and $ zipWith f (group s) (group t)
  where
    f a b = head a == head b && (la == 1 && lb == 1 || 1 < la && la <= lb)
      where
        la = length a
        lb = length b

問題が二つある。
ひとつは、String を長さ $2\times10^5$ に使うのは贅沢なこと。
もうひとつは、zipWith は片方のリストに余りが出たときに無視することである。
結局、ByteString を対象に、自前で再帰する計算で実現する必要がある。幸い、ByteStringにもgroup, length, head はある。

結果

import qualified Data.ByteString.Char8 as BS

abc259c :: BS.ByteString -> BS.ByteString -> Bool
abc259c s t = loop (BS.group s) (BS.group t)

loop (a:as) (b:bs) = f a b && loop as bs
loop [] [] = True
loop _  _  = False

f a b = BS.head a == BS.head b && (la == lb || 1 < la && la <= lb)
  where
    la = BS.length a
    lb = BS.length b

長さの条件を「両方の長さが等しいか、Sが2以上ならばTはそれ以上でも構わない（短いのはダメ）」と変えてみた。意味は前と同じ。

D - Circumferences

問題 ABC259D

シグネチャを決める。

abc259d :: Int             -- N
        -> Int -> Int      -- sx, sy
        -> Int -> Int      -- tx, ty
        -> [(Int,Int,Int)] -- xi,yi,riのタプルのリスト xyrs
        -> Bool

距離は基本的に2乗のまま扱うことにして、整数で話を進める。

ステップは次のように3つに分解できる。

1. 開始点 $(s_x, s_y)$、終着点 $(t_x, t_y)$ を円周に乗せているのは何番の円かを知る（複数の可能性もあるが、いずれか一つでよい）それぞれ $a, b$ 番とする
2. 二つの円が共有点を持つかを知る
3. $a$ 番から $b$ 番まで、共有点を持つ円を辿って到達できるかを知る

ステップ1は、それぞれの円の中心との距離が半径と等しいような点をひとつ見つければよい。

dist2 x1 y1 x2 y2 = (x1 - x2)^2 + (y1 - y2)^2

findab x y = head [i | (i, (xi, yi, ri)) <- zip [0..] xyrs, dist2 x y xi yi == r^2]

a = findab sx sy
b = findab tx ty 

ステップ2は、場合を分けて考える必要がある。
普通の、両方とも円の中心が相手の円の外にある場合、二つの半径の和と中心間の距離とが等しいとき接していて、より大きいときは離れていて、より小さいときは2点で交わっている。
普通でない、一方の円の中にもう一方の円の中心が入っている場合、半径の差と中心間の距離が等しいとき接していて、より小さいときは離れていて、より大きいときは2点で交わっている。

connect (x1,y1,r1) (x2,y2,r2)
  | (max r1 r2)^2 <= d2 = (r1 + r2)^2 >= d2
  | otherwise           = (r1 - r2)^2 <= d2
  where
    d2 = dist2 x1 y1 x2 y2

ステップ3はまず、ステップ2を使って $N$ 個の円について総当たりで共有点を持つ円の組み合わせを調べる。この組み合わせを伝って $a$ から $b$ に行けるかどうかは、無向グラフの辺の情報を与えられて、二つのノード間が連結であるかを判定する問題といえる。
実はこれは Union-Find もしくは Disjoint Set Union （素集合データ構造）というものを使うと効率的に判定できる定番の問題である。要素を0から上限までの整数に限定した簡単な実装を示す。これはData.Arrayの更新を使っているので効率はよくない。

type UnionFind = Array Int Int

-- 0からN-1までのN要素が独立した初期状態を作る
newUF :: Int -> UnionFind
newUF n = listArray (0, pred n) $ replicate n (-1)

-- 補助関数 ノードの根まで辿り、そのrank（負）と共に返す
getRoot :: UnionFind -> Int -> (Int, Int)
getRoot uf i = let k = uf ! i in if k < 0 then (i, k) else getRoot uf k

-- ふたつのノードが同じ分割に属しているか判定する
findUF :: UnionFind -> Int -> Int -> Bool
findUF uf i j = fst (getRoot uf i) == fst (getRoot uf j)

-- ふたつのノードが同じ分割に属していることを登録する
uniteUF :: UnionFind -> Int -> Int -> UnionFind
uniteUF uf i j
  | a == b = uf
  | otherwise =
      case compare r s of
        GT -> uf // [(a,b)]
        LT -> uf // [(b,a)]
        EQ -> uf // [(b,a), (a, pred r)]
  where
    (a, r) = getRoot uf i
    (b, s) = getRoot uf j

共有点を持つ円の番号どうしを uniteUF で全て登録し、最後に $a$ と $b$ で findUF することで、$a$ から $b$ まで辿れるかが判定できる。

uf = foldl (\uf (i,j) -> uniteUF uf i j) (newUF n)
  [ (i,j)
  | (i,c1) : jcs <- tails (zip [0..] xyrs)
  , (j,c2)       <- jcs
  , connect c1 c2]

findUF uf a b -- 答え

結果

import Data.Array
import Data.List

abc259d :: Int -> Int -> Int -> Int -> Int -> [(Int,Int,Int)] -> Bool
abc259d n sx sy tx ty xyrs = findUF uf a b
  where
    a = findab sx sy xyrs
    b = findab tx ty xyrs
    uf = foldl (\uf (i,j) -> uniteUF uf i j) (newUF n)
      [ (i,j)
      | (i,c1) : jcs <- tails (zip [0..] xyrs)
      , (j,c2)       <- jcs
      , connect c1 c2]

dist2 x1 y1 x2 y2 = (x1 - x2)^2 + (y1 - y2)^2

findab x y xyrs = head
  [i | (i, (xi, yi, ri)) <- zip [0..] xyrs, dist2 x y xi yi == ri^2]

connect (x1,y1,r1) (x2,y2,r2)
  | (max r1 r2)^2 <= d2 = (r1 + r2)^2 >= d2
  | otherwise           = (r1 - r2)^2 <= d2
  where
    d2 = dist2 x1 y1 x2 y2

-- UnionFindのコードは省略

Data.Array実装が不安だったがTLEにならずACすることを確認できた。

E - LCM on Whiteboard

問題 ABC259E

シグネチャを決める。$m_i$ は捨てる。

abc259e :: Int           -- N
        -> [[(Int,Int)]] -- pとeのペアのリストのリスト pess
        -> Int           -- 答え

問題は、整数 $N$ 個のうち、ひとつを除いた $N-1$ 個の最小公倍数(LCM)のバリエーションがいくつあるか、と尋ねている。
素因数分解で表現されている数のLCMとはつまり、素因数 $p$ ごとにべき乗している数 $e$ の最大値を選択したものである。
しかし、$p$ をキー、$e$ を値とする写像もしくは、$p$に背番号を振って $e$ の配列を作り、これをセグメント木にあてはめて、要素ごとにmaxをとる方法では時間がかかりすぎる。

それぞれの素因数 $p$ について、$N-1$ 個の候補 $e$ のうちいずれかが最大値として選ばれる。よく考えるとそのバリエーションは $N$ 通りもあるわけではない。例えば、$S = \{1,2,3,4,5\}$としたとき、$\max(S \setminus \{5\}) = 4$, $1$ から $4$ の $i$ についてはどれも $\max(S \setminus \{i\}) = 5$ で同じである。つまり、最大値と2番目に大きい値の2つだけが現れうる。

そこでまずは、現れる全ての $p$ に対して、$e$ の最大値と次点を見つける。ただし最大値が複数回出現したときは次点は最大値と同じ値とする。
降順ソートして先頭2つを取る、という無精をすると $O(N \log N)$ かかるといけないので、線形に $O(N)$ で済む形で実現する。
IntMap に最大値と次点のペアを対応付けて、最大値が更新される、次点が更新される、更新されない、という計算を累積する。
（fの第1引数のパターンで、ワイルドカードにせずあえて0を書いているのは、それがリストから与えられた新たな要素であることを確実にするため。fromListWithのドキュメントを見ても、2つの引数のどちらに新たな値、どちらに累積された値が来るかは、例から推測する形でしか書かれていない。）

-- p に対して (eの最大値, eの次点) を割り当てる写像
maxAnd2nds = IM.fromListWith f [(p,(e,0)) | pes <- pess, (p,e) <- pes]

f (e,0) (a,b)
  | a < e = (e,a)     -- 最大値更新
  | b < e = (a,e)     -- 次点更新
  | otherwise = (a,b) -- 更新なし

さて、maxAnd2nds が得られたら、これを利用して、数をひとつ除いたときのLCMがどうなるかを調べることができる。
ある数を候補から除外したとき、それぞれの $p$ について $e$ の最大値か次点がLCMに使われる。そこで、ある数を除外したときに最大値が使われなくなるような $p$ のリストで、その場合のLCMを特徴づけることができる。（逆に、最大値が使われる $p$ のリストで表してもできるが、pess からそれを作るのはより面倒である。）
ここで注意が必要なのは、最大値と次点が等しい場合（最大値が複数回出現している場合）には、$e$ が最大値で次点を選んでも結果は同じ値のままということである。（このため上で「次点が使われる」ではなく「最大値が使われなくなる」という言い回しにした。）
数を表す pes から、それを特徴づける $p$ のリストは次の関数で求められる。幸いにも $p$ は初めから昇順であるのでソートせずとも正規化できている。

feature pes = [p | (p,e) <- pes, let (a,b) = maxAnd2nds IM.! p, e == a, a /= b]

これのバリエーションを数えたら、求める答えが得られる。

S.size $ S.fromList $ map feature pess

結果

import qualified Data.IntMap as IM
import qualified Data.Set as S

abc259e :: Int -> [[(Int,Int)]] -> Int
abc259e n pess = S.size $ S.fromList $ map feature pess
  where
    maxAnd2nds = IM.fromListWith f [(p,(e,0)) | pes <- pess, (p,e) <- pes]
    feature pes = [p | (p,e) <- pes, let (a,b) = maxAnd2nds IM.! p, e == a, a /= b]

f (e,0) (a,b)
  | a < e = (e,a)
  | b < e = (a,e)
  | otherwise = (a,b)