入水目指して過去挑戦して解けてないAtCoder Problems上でDifficultyが水色の問題を解いていく。
ABC329 F: Colored Ball概要
N個の箱に1つずつ色のついたボール(色の重複はあり)が入っていて、指定された箱から全てのボールを別の箱に移動させて色の種類数を求めることをQ回実施する問題。
解法イメージ
愚直な方針
まず、問題文の通りに解くことを考えてみる。
N個の箱に1つずつボールが入っているので、1つのボールが入っている配列の配列を作成する。($ \mathcal{O}(N) $)
その後Q個のクエリを処理する。($ \mathcal{O}(N+Q) $)
各クエリに対して、箱aに入っているボール(最大N個)を全て箱bに移動させる。($ \mathcal{O}(N+QN) $)
移動させた後、箱bに入っているボールの色の種類数を求めるて出力する。($ \mathcal{O}(N+QN^2) $)
という感じで、そのまま解くと$ \mathcal{O}(QN^2) $となり、$ 1 \le N,Q \le 2 * 10^5 $の制約から当然間に合わない。
種類数の計算の効率化
計算量の削減を考えた時に、ボールを移動する際は全てのボールを移動させるためそれぞれの色について何個入っているかを考える必要はない。
よって、色の種類数のみを保持してあげれば良いので、Set等のオブジェクトで持ってあげてボールを移動させる際に色の種類が増えたかどうかを求めてその箱の種類数を管理してあげることで箱に入っている色の種類数を$ \mathcal{O}(1) $で求めることができる。(計算量$ \mathcal{O}(N+QN) $)
ボールの移動方法の効率化
あとはボールの移動の計算量が$ \mathcal{O}(QN) $になっている部分を改善できれば良さそう。
移動する時の回数がNになるような移動方法を考えてみると、
箱1→箱2
箱2→箱3
...
箱N-1→箱N
箱N→箱1
...(以下略)
のように一番ボールが多い箱の移動をしていくケースが考えられる。
このケースを考える時に、移動元の箱にはたくさんボールが入っているが、移動先の箱には1個しか入っていない状態になっている。
それなら、移動先のボールを移動元に移動させて、箱のラベルの張り替えをした方が移動回数は減らせそうに見える。
実際に上記のケースを考えると、箱N-1→箱Nの移動までは移動先のボール一つを移動元に移動させてラベルの張り替える2ステップで良く。それ以降は全てのボールが一つの箱に集まっているためラベルの張り替えのステップのみで済むため$ \mathcal{O}(Q) $となる。
これは効率のいいケースなので、一番移動回数が多くなりそうなケースを考えてみる。
ボールの移動する回数を増やすには、移動先と移動元のボールの数の小さい方の数を大きくすれば良いので、これを最大化させるには移動先と移動元のボールの数を同数にすることを考える。
初期状態では全て1つずつ箱に入っているため、それぞれ1つの箱から1つの箱に移動させてボールが2つの箱ができる。それらの箱に対して、2つずつ入っている箱同士で移動させ4つの箱を作る。というのを繰り返していくと、$ 2^k $個のボールの箱ができるまでに
$$ 1(個) \times \frac{N}{2}(箱) + 2(個) \times \frac{N}{4}(箱) + ... + 2^{K-1}(個) \times \frac{N}{2^K}(箱) = \frac{N}{2} + \frac{N}{2} + ... + \frac{N}{2} = \frac{KN}{2} $$
から$ \frac{KN}{2} $回のボールの移動が発生する。
この移動に必要なクエリの回数は、指定した箱数と一致するので$ \Sigma_{i=1}^{K}\frac{N}{2^{i}} = NK(1-\frac{1}{2^K}) $となる。
よって1回のクエリあたり $ 2 - \frac{1}{2^{K-1}} $個のボールが移動することになるので、こちらで考えても$ \mathcal{O}(Q) $で処理ができる。
※実際にはここまで考える必要はなくて、石が移動するクエリの回数は1回進めるごとに半減していくため、$ \mathcal{O}(\log{N}) $であり、それぞれの移動は最大でもN個なので$ \mathcal{O}(N\log{N}) $で済む。これだけでも、$ \mathcal{O}(Q+N\log{N}) $程度で収まることが分かるので間に合うことは分かる。
よって、ボールの移動は数が少ない箱から数が多い箱へ移動し、移動した箱を移動先の箱としてラベリングしてあげればよい。
実装例
require "set"
n, q = gets.split.map(&:to_i)
arr = gets.split.map(&:to_i)
hash = {}
arr.each_with_index do |a, i|
hash[i + 1] = Set.new([a])
end
q.times do
a, b = gets.split.map(&:to_i)
if hash[b].length > hash[a].length
hash[b].merge(hash[a])
else
hash[a].merge(hash[b])
hash[b] = hash[a]
end
hash[a] = Set.new
puts hash[b].length
end