백준 4195번 친구 네트워크

Description

민혁이는 소셜 네트워크 사이트에서 친구를 만드는 것을 좋아하는 친구이다. 우표를 모으는 취미가 있듯이, 민혁이는 소셜 네트워크 사이트에서 친구를 모으는 것이 취미이다.

어떤 사이트의 친구 관계가 생긴 순서대로 주어졌을 때, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오.

친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말한다.

Input

첫째 줄에 테스트 케이스의 개수가 주어진다. 각 테스트 케이스의 첫째 줄에는 친구 관계의 수 F가 주어지며, 이 값은 100,000을 넘지 않는다. 다음 F개의 줄에는 친구 관계가 생긴 순서대로 주어진다. 친구 관계는 두 사용자의 아이디로 이루어져 있으며, 알파벳 대문자 또는 소문자로만 이루어진 길이 20 이하의 문자열이다.

# 예제 입력
2
3
Fred Barney
Barney Betty
Betty Wilma
3
Fred Barney
Betty Wilma
Barney Betty

Output

친구 관계가 생길 때마다, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오.

#예제 출력
2
3
4
2
2
4

Solution

# 내 풀이
n = int(input())

for _ in range(n):
    graph = dict()
    key = 0
    m = int(input())
    for i in range(m):
        data = list(input().split())
        if data[0] not in graph:
            key += 1
            graph[data[0]] = key

        if data[1] not in graph:
            graph[data[1]] = graph[data[0]]
        else:
            a = graph[data[1]]
            for x in graph:
                if graph.get(x) == a:
                    graph[x] = graph[data[0]]

        index = graph[data[0]]
        count = 0
        for x in graph:
            if graph.get(x) == index:
                count += 1
        print(count)

파이참에서는 결과값이 잘 나왔지만, 3중 for문의 사용으로 시간 초과를 받았다.

Feedback

Union Find 알고리즘을 이용

  1. 노드들을 저장할 딕셔너리(parent)와 해당 노드의 수를 저장할 딕셔너리(number)를 선언한다.
  2. parent에 노드가 없다면 parent에 새로운 노드를 저장하고, number에 해당 노드의 수를 저장한다.
  3. x가 root노드인지 확인하고, root 노드가 아닐 경우 재귀 함수를 통해 root 노드를 찾는 find 함수를 선언한다.
  4. find 함수로 찾은 root 노드 값이 다를 경우 한 트리의 root 노드를 다른 트리의 root 노드의 자식으로 연결해주는 union 함수를 정의한다.
  5. union 함수를 통해 노드 사이를 이은 뒤, 해당 순회 때 x의 root 노드를 키 값으로 갖는 number의 value를 출력한다.
# 모범 풀이
def find(x):
    if x == parent[x]:
        return x
    else:
        p = find(parent[x])
        parent[x] = p
        return parent[x]

def union(x, y):
    x = find(x)
    y = find(y)
    if x != y:
        parent[y] = x
        number[x] += number[y]

test_case = int(input())

for _ in range(test_case):
    f = int(input())
    parent = dict()
    number = dict()

    for _ in range(f):
        x, y = input().split(' ')
        if x not in parent:
            parent[x] = x
            number[x] = 1
        if y not in parent:
            parent[y] = y
            number[y] = 1

        union(x, y)
        print(number[find(x)])

union find 알고리즘을 사용하여 트리와 같은 자료구조를 구현했더니 for문을 사용하여 네트워크가 생길 때마다 모든 value값을 업데이트 할 필요가 없어졌고, 2중 for문으로 문제를 해결할 수 있었다.

Source