paiza「【殿堂入りキャンペーン】村人の友好関係 (paizaランク S 相当)」をHaskellで

問題

制約を確認。

• $2 \leq N \leq 500$ ：村人は500人まで
• $1 \leq M \leq N(N-1)/2$ ：0でない友好度は総当たりの上限 $O(N^2)$ 個まであり
• $1 \leq Q \leq 50,000$ ：クエリは $5\times 10^4$ 回まで

シグネチャを決める。横着する。

solve :: [Int]         -- N,M,Q
      -> [[Int]]       -- a_i, b_i, f_i
      -> [(Bool, Int)] -- op_i (+のときTrue), q_i
      -> [Int]         -- 答え

自分の挑戦と敗北

特に上手いアルゴリズムが当てはまりそうもないので、状況を素朴にシミュレーションしてみる。

• 現在のメンバー一覧を IntSet で管理
• 現在、メンバーの境界を跨いでいる友好度の多重集合を IntMap で管理　キーが友好度、値が多重度、0になったとき、キーも削除することで、findMax 一発で最大値が得られる
• クエリに対して、入るなり出るなりしたその人と、他$N-1$人との間の友好度を、
  • 同じ位置になったのなら、友好度集合から除く
  • 違う位置になったのなら、友好度集合に加える
    という操作を行う。1人あたり $O(\log (N^2))$ なので全員分で１クエリあたり $O(2N \log N)$
  • findMax で現在の最大値を取り出す $O(2\log N)$
• をクエリの個数だけ繰り返す。

全体で $O(QN \log N)$ これは $5 \times 10^4 \cdot 500 \cdot \log 500 \fallingdotseq 2.5 \times 10^8$ 程度。少し大きいけどなんとか間に合うかな？

IntSetの読み出しに時間がかかるのを嫌って、メンバーは Array Int Bool で管理してみる。

import Data.Array
import qualified Data.IntMap as IM

type MultiIntSet = IM.IntMap Int

solve :: [Int] -> [[Int]] -> [(Bool,Int)] -> [Int]
solve [n,m,q] abfs qs = snd $ mapAccumL step (mem0, mis0) qs
  where
    f = accumArray (flip const) 0 ((1,1),(n,n)) $
        concat [[((a,b),f), ((b,a),f)] | a:b:f:_ <- abfs]
    mem0 = accumArray (||) False (1,n) []
    mis0 = singletonMIS 0 -- 0は番兵
    step (mem, mis) (op, r) = ((mem1, mis1), findMaxMIS mis1)
      where
        mem1 = mem // [(r, op)]
        mis1 = foldl' misf mis (assocs mem)
        misf mis (s, b)
          | r == s    = mis -- 自分自身について、何もしない
          | op == b   = deleteMIS (f ! (r,s)) mis -- rの移動先と同じ状態なら、外す
          | otherwise = insertMIS (f ! (r,s)) mis -- rと分かれる相手の数字は、追加する

singletonMIS x = IM.singleton x 1

insertMIS x ms = IM.insertWith (+) x 1 ms

deleteMIS x ms = IM.update dec x ms
  where
    dec 1 = Nothing
    dec n = Just $ pred n

findMaxMIS ms = fst $ IM.findMax ms

結果、テスト6が3秒経過でタイムアウト。他は何ともない。

この後、mutableな配列にしてみたり色々と係数を稼ぐ細工をしてみたが、効果なし。どうなってんの？

公式解説

他の問題は解説やコード例がついてるのに、これは「解答コード例・解説なし」えっナンデ？

Qiitaの記事

この問題を扱っている他の記事に頼ろう。

A: numPyパワー

メンバー表を Bool の配列 g に、友好度を二次元配列 m[] に入れておくと、m[g][:, ~g] で「gが真な添え字を1つめの添え字に、偽な添え字を2つめの添え字にした部分配列を取り出せちゃうよ」ということなのかな？
力任せにやってるだけにも程があるけど、numPyならこれで間に合うんだ。何それ。

友好度の配列の探す範囲は最大でその面積の1/4で、$O(N^2/4)$ これをクエリの回数だけ繰り返すので全体では $O(QN^2)$ これは $O(QN \log N)$ より大きいんだけどなぁ。

B: 友好度を線形探索

ついつい、友好度を二次元配列に納めたくなるのを改め、友好度の大きい順にソートしてリストとしてとっておく。
入退会が起きた後で、このリストを前から舐めて、中と外に分かれているような、最初に見つかった関係が答え。

リストの長さは$M$なので線形探索の時間は$O(N^2)$で、全体ではやはり $O(QN^2)$

C: 調べるべき友好度をもっと絞る

クエリに対する仕事はBと同じだが、線形探索する友好度リストを、ただ大きい順で並べ替えるだけでなく、不要なものを取り除くことで数を減らしている。

そのために、友好度の大きい方を優先する最(大)スパニング木の辺になる友好関係だけを残す。友好度の低い、ショートカットする辺は、より友好度の高い辺に上書きされるので、決して使われないから、というのがミソ。
こうすると、最大スパニング木の森にある辺だけが線形探索の対象になるので、要素数は $N-1$ が上限。つまり計算量は$O(QN)$ となる。これが正解か…

D: C++パワー

これは自分の敗退解と同じアプローチで、map <long long, int> scores がmultisetを表現している。

ただ最後の行

for(;scores.size()>2000;)scores.erase(scores.begin());

「multisetの要素が2000を越えないように、小さい方から削る」が何でそんなことしちゃっていいのか全くわけわからん。
「こうしないと要素数が多すぎてTLEする、これより削ると消しちゃ駄目な要素を消してWAする」んじゃないだろな…と思って試したところ、どうやらそうみたい。これ嘘解答じゃないですかヤダー！

ていうかこのアプローチ、C++で書いても間に合わんのな…

再挑戦

Aメソッドは、素のHaskellで書いてもまぁ無理だろうからやめておく。

ならまず簡単なBの方で。

import qualified Data.IntSet as IS

solve :: [Int] -> [[Int]] -> [(Bool,Int)] -> [Int]
solve [n,_m,_q] abfs qs = snd $ mapAccumL step IS.empty qs
  where
    abfs1 = sortBy cmp abfs
    cmp (_:_:f1:_) (_:_:f2:_) = compare f2 f1
    step mem (b,r) = (mem1, ans)
      where
        mem1 = (if b then IS.insert else IS.delete) r mem
        ans = head $ [f | a:b:f:_ <- abfs1, IS.member a mem1 /= IS.member b mem1] ++ [0]

テスト6: 0.21sec, テスト7: 0.50sec で通ってしまった。あっさり。
あまりに速いので、ちゃんとしたCメソッドをする元気が出ないよ。

Bで十分間に合ってしまったのは、実用上の傾向なのか、単なるテストケース不足なのか。
それとも計算機が速くなった分？

まとめ

この問題設定において、MSTを使うことで、調べる関係の個数を $N(N-1)/2$ から $N-1$ に減らすことができるというのは重要で、本当はそうやるべきなので、次に必要になる場面のために覚えておこう。

終わりに

スキルチェックでHaskell使えるようにしてホスィ…