[BOJ] 4195 친구 네트워크

유니온-파인드


이 문제는 친구를 모아서 친구의 수를 세라고 한다. 딱 보면 유니온-파인드 써서 같은 숫자를 갖는 원소의 개수를 어떻게 셀 것인가에 대한 내용이다.

그래서 우선 유니온 파인드를 어떻게 하면 될지 생각을 해보자. 우선 그냥 parent 배열을 순회하며 세면 O(10만)이고, 이를 10만번..? 그리고 여러 테스트케이스까지 한다면 시간초과가 날 것이 분명하다. 따라서 어떻게 순회를 하지 않고 count를 세어줄 수 있을까? 해서 배열을 하나 더 도입했다. union을 할 때마다 양쪽 root 인덱스에 count값을 각각 더해서 반환한다. 자세한 알고리즘은 다음과 같다.

  1. 친구 이름이 영어로 들어오니 Key: Value로 이름: 숫자 형식으로 dictionary를 이용해 맵핑을 해주자.
  2. 부모를 가리키는 parent 배열을 모두 자기 자신을 가리키도록 초기화한다.
  3. 네트워크에 속한 친구의 수를 세어줄 count 배열을 모두 초기화한다.
  4. 전체 connection 정보를 순회하며 union 연산
    1. 만약 이미 두 친구가 같은 집합이라면 카운트를 더하지 않음
    2. 다른 집합이라면 더 작은 루트노드로 몰아준다. 카운트는 둘다 더해줌.

find 연산은 경로압축법을 이용해 시간복잡도를 개선했다.

구현


import sys

def find(x, parent):
    if x != parent[x]:
        parent[x] = find(parent[x], parent)
    return parent[x]

def union(a, b, parent, count):
    a = find(a, parent)
    b = find(b, parent)
    answer = count[a] + count[b]
    if a < b:
        parent[b] = a
        count[a] = count[b] = answer
    elif a > b:
        parent[a] = b
        count[a] = count[b] = answer
    else:
        answer = count[a]
    return parent, count, answer

def solution(n):
    parent = [i for i in range(n)]
    count = [1 for _ in range(n)]
    for num1, num2 in connections:
        parent, count, answer = union(num1, num2, parent, count)
        print(answer)

if __name__ == "__main__":
    T = int(sys.stdin.readline().rstrip())
    for _ in range(T):
        F = int(sys.stdin.readline().rstrip())
        idx = 0
        connections = []
        name_to_num = dict()
        for _ in range(F):
            name1, name2 = map(str, sys.stdin.readline().split())
            if name1 not in name_to_num:
                name_to_num[name1] = idx
                idx += 1
            if name2 not in name_to_num:
                name_to_num[name2] = idx
                idx += 1
            connections.append([name_to_num[name1], name_to_num[name2]])
        solution(len(name_to_num))

태그:

카테고리:

업데이트: