🌳

[ABC049 D] 連結

2024/01/27に公開

Ruby

二つの森を扱う難問。

問題文を理解する

きっちり書かれた文章がなかなか頭に入ってこない病なので、

全ての都市について、その都市と道路・鉄道のどちらでも連結している都市の数を求めてください。

を、

・それぞれの街から見て道路と鉄道の両方で行ける街は？

に言い換える。

解法

街Xから順に街 1, 2, 3, 4 に行くとする
そのとき「道路と鉄道の両方」を使って行けるか試す
一通り試したら行けた回数を表示する

これを街1から街4まで繰り返す。

例1のシミュレーション

道路

すべて繋がっている。

鉄道

一部だけが繋がっている。

疑問

ところで同じ街には鉄道で行けるのか？たとえば街1から街1に行けるのか？現実的に考えたら同じ街の中で鉄道を敷く意味がないので鉄道は通っていないと考えてよさそうだが、

すべての都市はそれ自身と道路で連結しているとみなします。鉄道についても同様に定めます。

と問題文に書いてあるので通っている。このように現実的にはちょっと変だけど問題にとって都合のよい設定になっていたりするので気にかけておく。

街1基点

そこで街1を基点にして街 1, 2, 3, 4 に行くとしたとき、

移動先	道路	鉄道
1	○	○
2	○
3	○
4	○

道路と鉄道の両方で行けるのは、両方○なので 1 つだとわかる。

同様にして、街2〜4も調べると

街 2, 3, 4 基点

移動先	道路	鉄道
1	○
2	○	○
3	○	○
4	○

移動先	道路	鉄道
1	○
2	○	○
3	○	○
4	○

移動先	道路	鉄道
1	○
2	○
3	○
4	○	○

となるため、答えは 1, 2, 2, 1 になる。

コード (TLE)

上のシミュレーションの通りに組むと次のようになる。

N, K, L = gets.split.collect(&:to_i)                 # => [4, 3, 1]
Ks = K.times.collect { gets.split.collect(&:to_i) }  # => [[1, 2], [2, 3], [3, 4]]
Ls = L.times.collect { gets.split.collect(&:to_i) }  # => [[2, 3]]
Nodes = 1..N
ds1 = DisjointSet[Ks]           # 道路の森
ds2 = DisjointSet[Ls]           # 鉄道の森
counts = Hash.new(0)
Nodes.each do |a|
  Nodes.each do |b|
    if ds1.same?(a, b) && ds2.same?(a, b)
      counts[a] += 1
    end
  end
end
counts.values                                        # => [1, 2, 2, 1]
puts counts.values * " "

ところが提出すると TLE になってしまう。

普通の問題では、明示していない素集合データ構造に気づくかどうかが鍵になっているのだが、この問題は、その先に真のボスが待っていた。

TLE の原因

Nodes.each {} の二重ループが原因。そこが O(n^2) になっているので、same? を経路圧縮で高速化してもそもそも回しすぎなのだった。

最適化

道路と鉄道の根のペアの出現頻度を調べる。

根のペアから、

道路の根	鉄道の根
2	1
2	3
2	3
2	4

頻度に変換し、

道路の根	鉄道の根	頻度
2	1	1
2	3	2
2	4	1

答えは街1から順に並べるので再度の根のペアをキーにして引く。これで O(n^2) から O(2n) になって通る。

コード (AC)

N, K, L = gets.split.collect(&:to_i)                          # => [4, 3, 1]
Ks = K.times.collect { gets.split.collect(&:to_i) }           # => [[1, 2], [2, 3], [3, 4]]
Ls = L.times.collect { gets.split.collect(&:to_i) }           # => [[2, 3]]
ds1 = DisjointSet[Ks]                                         # => 親:{1=>2, 2=>2, 3=>2, 4=>2} 構図:{2=>[1, 2, 3, 4]} 節数:{2=>4} 木数:1 高さ:{2=>2}
ds2 = DisjointSet[Ls]                                         # => 親:{2=>3, 3=>3} 構図:{3=>[2, 3]} 節数:{3=>2} 木数:1 高さ:{3=>2}
Nodes = 1..N
roots = Nodes.collect { |e| [ds1.root(e), ds2.root(e)] }      # => [[2, 1], [2, 3], [2, 3], [2, 4]]
freq = roots.tally                                            # => {[2, 1]=>1, [2, 3]=>2, [2, 4]=>1}
ans = Nodes.collect { |e| freq[[ds1.root(e), ds2.root(e)]] }  # => [1, 2, 2, 1]
puts ans * " "

最適化問題

最後の最適化は慣れていないと相当難しい。そこだけ単純化して別の問題にするとこんな感じだろうか。

a b の順番で出現頻度を求めよ

Chars = ["b", "b", "a"]

# 遅い O(n^2)
counts = Hash.new(0)
["a", "b"].each do |x|
  Chars.each do |y|
    if x == y
      counts[x] += 1
    end
  end
end
counts.values                       # => [1, 2]

# 速い O(2n)
freq = Chars.tally                  # => {"b"=>2, "a"=>1}
["a", "b"].collect { |e| freq[e] }  # => [1, 2]