[ BOJ / Python ] 4195번 친구 네트워크

황승환·2022년 8월 16일
0

Python

목록 보기
445/498

이번 문제는 유니온-파인드 알고리즘을 통해 해결하였다. 기존 유니온-파인드와 다르게 값이 str형식으로 들어왔기 때문에 딕셔너리를 사용하였다. 처음에는 딕셔너리를 통해 부모를 찾고, 아직 못찾았을 부모가 있을 수 있으므로 부모 딕셔너리를 순회하며 find함수를 다시 돌려 그 값이 현재 입력된 값의 부모와 같다면 값을 더하도록 하였다. 그러나 이 방식에서는 시간초과가 발생하였다.

그래서 부모 딕셔너리와 같은 네트워크 내의 수를 세는 딕셔너리를 이용하였다. union함수에서 부모를 지정할 때 네트워크 내의 수를 더하도록 하여 굳이 부모 딕셔너리를 순회하지 않고도 바로 값을 도출할 수 있도록 하였다.

Code

처음 코드 (시간 초과)

from collections import defaultdict
t = int(input())
def find(a):
    if a != parents[a]:
        parents[a] = find(parents[a])
    return parents[a]
def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        tmp = sorted([a, b])
        parents[tmp[1]] = tmp[0]
for _ in range(t):
    f = int(input())
    parents = defaultdict(str)
    for _ in range(f):
        a, b = map(str, input().split())
        if not parents[a]:
            parents[a] = a
        if not parents[b]:
            parents[b] = b
        union(a, b)
        answer = 0
        for key, value in parents.items():
            if find(value) == parents[a]:
                answer += 1
        print(answer)

정답 코드

from collections import defaultdict
t = int(input())
def find(a):
    if a != parents[a]:
        parents[a] = find(parents[a])
    return parents[a]
def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        tmp = sorted([a, b])
        parents[tmp[1]] = tmp[0]
        nums[tmp[0]] += nums[tmp[1]]
for _ in range(t):
    f = int(input())
    parents = defaultdict(str)
    nums = defaultdict(int)
    for _ in range(f):
        a, b = map(str, input().split())
        if not parents[a]:
            parents[a] = a
            nums[a] = 1
        if not parents[b]:
            parents[b] = b
            nums[b] = 1
        union(a, b)
        print(nums[find(a)])

profile
꾸준함을 꿈꾸는 SW 전공 학부생의 개발 일기

0개의 댓글