Zenn
🌳

[ABC231 D] Neighbors

2024/01/26に公開
Ruby
AtCoder
競プロ
素集合データ構造
disjointset
tech

https://atcoder.jp/contests/abc231/tasks/abc231_d

節が同じ木かを調べる問題。

シミュレーション

Bさんは人気者

  1. 机が横一列に並んでいる
  2. そこに A B C さんたちがやってきた
  3. B さんの隣にみんな座りたがる
  4. A さんは B さんの隣に座りたかった
    → A さんは B さんの左に座ったので A-B
  5. C さんも B さんの隣に座りたかった
    → C さんは B さんの右に座ったので A-B-C

となって (A B) (C B) は問題なし。

Dさん登場

そこに登場した D さんも B さんの隣に座りたいと言いだした。が、B さんの左右は埋まっているので座れなかった。これは、B さんを「3回以上指名」したときに、座れない人が出てくるということでもある。したがって (A B) (C B) のあとの (D B) は B が3回出てくるのでダメ。

移り気なAさん

また A さんは C さんの隣にも座りたいと言い出した。しかし、机は横一列に並べるので A-B-C の端と端を繋げることはできない。仮に A-B-C-D-E だったとしたときも (B D) を繋げると横一列にならなくなる。つまり端と端というより同じグループ内では隣同士にできない。したがって (A B) (C B) のあとの (A C) もダメ。

エッジケース

ここで (A B) のペアが再度来た場合が気になるが、そもそもペアはユニークだと記載があるので無視していい。もし来たとしてもすでに隣り合わせなので繋ぐことができないと考えていい。

解法

まとめると次の二つの条件を判定する。

  • 条件1: 3回指名された人がいる？
  • 条件2: 同じグループ内で繋げようとしている？

どちらにもマッチしなければ Yes になる。

条件1: 3回指名された人がいる？

を判定するのは、

av = [["A", "B"], ["C", "B"], ["D", "B"]]
av.flatten.tally.any? { |_, count| count >= 3 }  # => true

または、

(av.flatten.tally.values.max || 0) >= 3  # => true

でよい。

条件2: 同じグループ内で繋げようとしている？

を判定するのは、

ds = DisjointSet.new
ds.unite("A", "B")
ds.same?("A", "C")  # => false
ds.unite("C", "B")
ds.same?("A", "C")  # => true

でよい。ここで素集合データ構造が役に立つ。

最終的なコード

N, M = gets.split.collect(&:to_i)                    # => [4, 2]
AB = M.times.collect { gets.split.collect(&:to_i) }  # => [[1, 3], [2, 3]]
success = true
if (AB.flatten.tally.values.max || 0) >= 3
  success = false
else
  ds = DisjointSet.new
  AB.each do |edge|
    if ds.same?(*edge)
      success = false
      break
    end
    ds.unite(*edge)
  end
end
success                                              # => true
puts success ? "Yes" : "No"

条件 1, 2 の処理は疎結合にしたい。

参考

Discussion