[BOJ] 4195 친구 네트워크
유니온-파인드
이 문제는 친구를 모아서 친구의 수를 세라고 한다. 딱 보면 유니온-파인드 써서 같은 숫자를 갖는 원소의 개수를 어떻게 셀 것인가에 대한 내용이다.
그래서 우선 유니온 파인드를 어떻게 하면 될지 생각을 해보자. 우선 그냥 parent 배열을 순회하며 세면 O(10만)이고, 이를 10만번..? 그리고 여러 테스트케이스까지 한다면 시간초과가 날 것이 분명하다. 따라서 어떻게 순회를 하지 않고 count를 세어줄 수 있을까? 해서 배열을 하나 더 도입했다. union을 할 때마다 양쪽 root 인덱스에 count값을 각각 더해서 반환한다. 자세한 알고리즘은 다음과 같다.
- 친구 이름이 영어로 들어오니 Key: Value로 이름: 숫자 형식으로 dictionary를 이용해 맵핑을 해주자.
- 부모를 가리키는 parent 배열을 모두 자기 자신을 가리키도록 초기화한다.
- 네트워크에 속한 친구의 수를 세어줄 count 배열을 모두 초기화한다.
- 전체 connection 정보를 순회하며 union 연산
- 만약 이미 두 친구가 같은 집합이라면 카운트를 더하지 않음
- 다른 집합이라면 더 작은 루트노드로 몰아준다. 카운트는 둘다 더해줌.
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))