7K12 blog

猫でも分かるアルゴリズム解説

ABC157 D - Friend Suggestions

f:id:tkr987:20200708224848p:plain

https://atcoder.jp/contests/abc157/submissions/13265285

  • 芋づる式の構造は Union-Find (disjoint-set) にするのが典型テクニック
  • union-findと「二次元」配列を組み合わせて計算量を落とす
  • 二次元配列で細かく区切ることで無駄を無くす
  • i番目の人に対する友達候補の人数 = i番目の人と繋がりを持つ集合全体(友達の友達) - 直接的な友達の人数 - 自分自身 - (i番目の人と繋がりを持つ集合に含まれる)ブロック人数

f:id:tkr987:20200708234430p:plain
NもMもKも 10^5あることから O(N^2)で処理するとTLEしてしまう。4番目の条件が分かりづらいが、要するに友達の友達=友達候補ということ。芋づる式の構造になっているので、Union-Findを使って友達同士の集合を管理することで計算量を O(\log M)に改善する。

ただし、2番目の条件から直接友達同士になっている人を減算する必要がある。この処理はUnion-Findだけで実現しようとせず、二次元配列を併用することを考える。二次元配列を併用することで、i番目の人と直接友達になっている人数をfriend[i].size()でO(1)取得できるようになる。

さらに、1番目の条件からブロックリストになっている人も減算する必要がある。ただし、i番目の人とブロック関係の人数のうち「友達の友達」集合に含まれている人数だけを減算するようにしなければならない。これも二次元配列を利用して下記のように処理する。

rep(i, 0, N)

{ // i番目の人について友達候補を計算する

 ll ans = uf.Size(i) - 1 - Size(nyaaFriend[i]);

 for (auto&& e : nyaaBlock[i]) if (uf.Same(i, e)) ans--;

 cout << ans << " ";

}

二重ループで内側が線形処理になっているため O(NK)に見えるが、ブロックリストの合計 10^5しかないのが大事なところで、i番目の人の各ブロックリストサイズは小さい。そこで、二次元配列で細かく区切ってi番目の人関連以外はチェックしないようにすることで計算量を O(N+K \ log M)に抑えることが出来る。