🌳

[ABC157 D] Friend Suggestions

2024/02/22に公開

https://atcoder.jp/contests/abc157/tasks/abc157_d

SNSでの「おすすめユーザー」やECサイトでの「おすすめ商品」の出し方の基礎が学べる良問。

問題を理解する

問題の読み取りが難しいが、ようするに「各ユーザーに対するおすすめユーザー (の人数) を求めよ」ということらしい。

友達候補の定義を要約する

  • 候補者は自分ではない
  • 候補者をブロックしていない
  • 候補者をまだフォローしていない
  • 候補者は友達の友達(の友達……)である

要約してみると常識的なルールを言っているだけだったので気にかけなくていい。

解き方

  1. 全体の友達関係で巨大な素集合森を作る
  2. 各ユーザー毎に「自分の所属する木の大きさ」を求める
  3. 自分を除外する
  4. すでにフォローしていたら除外する
  5. ブロックしている人も除外する
  • (2) まで求めたらもう解けたようなもの

例1でシミュレーション

(1) 全体の友達関係で巨大な素集合森を作る。

と、なんか 1 を中心にしたネットワークに見えてしまうが、そこは繋っているかどうかだけが重要なので 2 3 4 から見ても対等である。したがって「自分の所属する木の大きさ」は全員 4 になる。あとは順に減らしていくだけ。

1 2 3 4 ペア 実行内容
4 4 4 4 (2) 各ユーザー毎に「自分の所属する木の大きさ」を求める
3 3 3 3 (3) 自分を除外する
(4) すでにフォローしていたら除外する
2 2 2-1
1 2 1-3
1 1 3-2
0 2 3-4
(5) ブロックしている人も除外する
0 1 4-1
0 1 0 1 最終的なおすすめ人数

コード (AC)

N, M, K = gets.split.collect(&:to_i)                 # => [4, 4, 1]
AB = M.times.collect { gets.split.collect(&:to_i) }  # => [[2, 1], [1, 3], [3, 2], [3, 4]]
CD = K.times.collect { gets.split.collect(&:to_i) }  # => [[4, 1]]

# (1) 全体の友達関係で巨大な素集合森を作る
ds = DisjointSet.unites(AB)                          # => 親:{2=>1, 1=>1, 3=>1, 4=>1} 構図:{1=>[2, 1, 3, 4]} 節数:{1=>4} 木数:1 高さ:{1=>2}

# (2) 各ユーザー毎に「自分の所属する木の大きさ」を求める
counts = {}
(1..N).each { |e| counts[e] = ds.size(e) }
counts                                               # => {1=>4, 2=>4, 3=>4, 4=>4}

# (3) 自分を除外する
counts.transform_values!(&:pred)

# (4) すでにフォローしていたら除外する
AB.each do |a, b|
  counts[a] -= 1
  counts[b] -= 1
end

# (5) ブロックしている人も除外する
CD.each do |c, d|
  if ds.same?(c, d)
    counts[c] -= 1
    counts[d] -= 1
  end
end

ans = counts.values                                  # => [0, 1, 0, 1]
puts ans * " "

Ruby っぽい書き方に注意する

各ユーザー毎に「自分の所属する木の大きさ」を求めるときシンプルに、

counts = (1..N).inject({}) { |a, e| a.merge(e => ds.size(e)) }  # => {1=>4, 2=>4, 3=>4, 4=>4}

と書いてしまうと TLE になる。冗長だけど、

counts = {}
(1..N).each { |e| counts[e] = ds.size(e) }
counts  # => {1=>4, 2=>4, 3=>4, 4=>4}

としないと通らない。これは inject が悪いのではなく merge で新しいハッシュを構築し続ける処理が最大10万回呼ばれるのがよくないんだと思われる。

意味のなかった最適化

すでにフォローしていたら除外する処理でデクリメントしているのが遅そうなので次のようにしたくなるが、

friends = Hash.new { |h, k| h[k] = [] }
AB.each do |a, b|
  friends[a] << b
  friends[b] << a
end
friends.each do |x, x_friends|
  counts[x] -= x_friends.count
end

AB のペアはユニークなのでとくに効果はない。それどころか O(M) でよかったのが O(M + N) になってしまう。

参考

Discussion