ABC157 D - Friend Suggestions

問題原文

atcoder.jp

問題要旨

SNSの友達関係・ブロック関係が無向グラフの形で与えられる。（友達関係の数は 個、ブロック関係の数は 個である）
ある人 について、友達関係のみを辿ってたどり着ける かつ その人とはブロック関係にない、という時、その人を「友達候補」と呼ぶ。
各人について友達候補の人数を求めよ。

• 
• 
• 
• 友達関係かつブロック関係となるような関係は与えられない

解法

各人について、「友達関係を辿ることで到達できるグループの大きさ - 自分の友達の数 - 自分の友達候補だがブロック関係の人の数 - １（自分自身）」 が答えになる。
問題は、「友達関係を辿ることで到達できるグループ」をどうやって得るか。これは UnionFind Treeか、DFS or BFSでの探索で求められる。

実装１：UF木

そしてこれは僕のUF木の記事。この記事での実装とはちょっぴり違うかもしれない。 at274.hatenablog.com

N, M, K = map(int, input().split())
F = [[] for _ in range(N)]
B = [[] for _ in range(N)]

for _ in range(M):
    a, b = map(int, input().split())
    a, b = a - 1, b - 1
    F[a].append(b)
    F[b].append(a)

for _ in range(K):
    c, d = map(int, input().split())
    c, d = c - 1, d - 1
    B[c].append(d)
    B[d].append(c)


class UnionFind:
    def __init__(self, n):
        self.par = [i for i in range(n)]
        self.rank = [0] * n
        self.size = [1] * n

    # 検索
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # 併合
    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
            self.size[y] += self.size[x]
        else:
            self.par[y] = x
            self.size[x] += self.size[y]
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # 同じ集合に属するか判定
    def same(self, x, y):
        return self.find(x) == self.find(y)

    # すべての頂点に対して親を検索する
    def all_find(self):
        for n in range(len(self.par)):
            self.find(n)


UF = UnionFind(N)
for iam in range(N):
    for friend in F[iam]:
        UF.union(iam, friend)

ans = [UF.size[UF.find(iam)] - 1 for iam in range(N)]
for iam in range(N):
    ans[iam] -= len(F[iam])

for iam in range(N):
    for block in B[iam]:
        if UF.same(iam, block):
            ans[iam] -= 1

print(*ans, sep=' ')

実装２：DFS

集合演算でゴリ押そうとするとTLEするので丁寧にやる。 （set - set の計算量は O(集合の大きさ) なので。）
どっちかというとオンラインクエリではないので、UFよりはこっちの実装の方が用途にぴったり合ってる気がする。
まぁ両方かけるべきとは思います。

N, M, K = map(int, input().split())
F = [[] for _ in range(N)]
B = [[] for _ in range(N)]

for _ in range(M):
    a, b = map(int, input().split())
    a, b = a - 1, b - 1
    F[a].append(b)
    F[b].append(a)

for _ in range(K):
    c, d = map(int, input().split())
    c, d = c - 1, d - 1
    B[c].append(d)
    B[d].append(c)


D = {}
parent = [-1] * N
visited = [False] * N
for root in range(N):
    if visited[root]:
        continue

    D[root] = set([root])
    stack = [root]
    while stack:
        n = stack.pop()
        visited[n] = True
        parent[n] = root

        for to in F[n]:
            if visited[to]:
                continue
            D[root].add(to)
            stack.append(to)

ans = [0] * N
for iam in range(N):
    group = D[parent[iam]]
    tmp_ans = len(group) - len(F[iam]) - 1
    for block in B[iam]:
        if block in group:
            tmp_ans -= 1
    ans[iam] = tmp_ans

print(*ans, sep=' ')

感想

本番は集合演算でゴリ押そうとして、無駄にペナを生やした。。。
UF典型といえば典型だけど、オンラインクエリでないからDFSの方がやっぱり肌にあってる気がする。