AtCoder Beginner Contest 157 D - Friend Suggestions (Python)

本日はこちらの水diff！！

問題はこちら〜

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

友達関係の問題！
現代社会にまつわる問題ですね（笑）

1. 問題の概要
2. 問題の制約と入力例
- 2-1. 問題の制約
- 2-2. 入力例
3. 問題を見て考えたこと
4. 解説
5. ACコード
6. 最後に

1. 問題の概要

N人の人々の間にM組の「友達関係」とK組の「ブロック関係」があります。

各人々 i に対して、友達候補の数を出力せよという問題。

ここでいう友達候補というのは簡単にいうとAさんから友達を伝って行ってBさんまでたどれればAとBは友達候補になります。

2. 問題の制約と入力例

2-1. 問題の制約

入力は全て整数
$\displaystyle 2 \leqq N \leqq 10^{5}$
$\displaystyle 0 \leqq M \leqq 10^{5}$
$\displaystyle 0 \leqq K \leqq 10^{5}$
$\displaystyle 1 \leqq A_{i}, B_{i} \leqq N$
$\displaystyle A_{i} \neq B_{i}$
$\displaystyle 1 \leqq C_{i}, D_{i} \leqq N$
$\displaystyle C_{i} \neq D_{i}$
$\displaystyle (A_{i}, B_{i}) \neq (A_{j}, B_{j}) \quad (i \neq j)$
$\displaystyle (A_{i}, B_{i}) \neq (B_{j}, A_{j})$
$\displaystyle (C_{i}, D_{i}) \neq (C_{j}, D_{j}) \quad (i \neq j)$
$\displaystyle (C_{i}, D_{i}) \neq (D_{j}, C_{j})$
$\displaystyle (A_{i}, B_{i}) \neq (C_{j}, D_{j})$
$\displaystyle (A_{i}, B_{i}) \neq (D_{j}, C_{j})$

2-2. 入力例

入力例1

$\displaystyle 4 \quad 4 \quad 1$
$\displaystyle 2 \quad 1$
$\displaystyle 1 \quad 3$
$\displaystyle 3 \quad 2$
$\displaystyle 3 \quad 4$
$\displaystyle 4 \quad 1$

出力例1

$\displaystyle 0 \quad 1 \quad 0 \quad 1$

入力例2

$\displaystyle 5 \quad 10 \quad 0$
$\displaystyle 1 \quad 2$
$\displaystyle 1 \quad 3$
$\displaystyle 1 \quad 4$
$\displaystyle 1 \quad 5$
$\displaystyle 3 \quad 2$
$\displaystyle 2 \quad 4$
$\displaystyle 2 \quad 5$
$\displaystyle 4 \quad 3$
$\displaystyle 5 \quad 3$
$\displaystyle 4 \quad 5$

出力例2

$\displaystyle 0 \quad 0 \quad 0 \quad 0 \quad 0$

3. 問題を見て考えたこと

友達の関係と言われたらぱっと思いつくのはやはり Union-Find ですよね。
でもどのように管理すれば良いかがよくわからない。

とりあえず「友達関係」の集合を Union-Find で管理してその後は...と言う所で少し躓いてしまいました。

それでは解説していきます。

4. 解説

この問題はお気づきの通り、 Union-Findで解くことができます。

まず、「友達関係」の集合を Union-Find で管理します。

その中で「友達候補」になることができないのは、自分自身もしくは「ブロック関係」にある人々です。

したがって、まずM個の入力で「友達関係」の集合を管理し、K個の入力で $\displaystyle C_{i}, D_{i}$ が同じグループに属していれば「友達候補」になれない集合（今回の場合は out 配列）に加えます。

そうすると、各人々 i に対して出力する友達候補の数は

$\displaystyle uf.size(i) - len(out[i]) - 1$

となります。
ここで out[i] := 人 i に対して友達候補になれない人数です。
$-1$ をしているのは自分自身を除くためです。

このようにしてこの問題は解くことが出来ました。

5. ACコード

class UnionFind:
  def __init__(self, n):
    self.n = n
    self.p = [-1] * n

  def leader(self, a):
    while self.p[a] >= 0:
      a = self.p[a]
    return a

  def merge(self, a, b):
    x = self.leader(a)
    y = self.leader(b)
    if x == y: return x
    if self.p[x] > self.p[y]:
      x, y = y, x
    self.p[x] += self.p[y]
    self.p[y] = x
    return x

  def same(self, a, b): return self.leader(a) == self.leader(b)

  def size(self, a): return -self.p[self.leader(a)]

n, m, k = map(int,input().split())

uf = UnionFind(n)

out = [[] for _ in range(n)] # 友達もしくは同じグループの中で喧嘩の組み合わせの集合
for _ in range(m):
    a, b = map(int,input().split())
    a -= 1
    b -= 1
    uf.merge(a, b)
    out[a].append(b)
    out[b].append(a)

for _ in range(k):
    c, d = map(int,input().split())
    c -= 1
    d -= 1
    if uf.same(c, d):
        out[c].append(d)
        out[d].append(c)

ans = []
for i in range(n):
    ans.append(uf.size(i) - len(out[i]) - 1)

print(*ans)

6. 最後に

ここ２日間ですが、早朝に起きて精進をする生活をしています。
朝を有効活用するととても得した気分になるので今後も続けて行けたらやっていきたいなと考えています。

Twitterのスペースを活用しているのでよければ作業がてら来てください！

最後まで見ていただきありがとうございました！