[백준 4195] 친구 네트워크

문제 출처 : https://www.acmicpc.net/problem/4195

4195번: 친구 네트워크

간단한 유니온 파인드 문제 입니다. 기존의 유니온 파인드와 다른 두 가지 때문에 쉽게 문제를 풀 수 없습니다. 그럼 문제점을 파악해보고 문제를 풀어보도록 하겠습니다.

문제점

이 문제의 가장 큰 문제는 노드가 숫자가 아닌 친구 이름이라는 것입니다. 해결할 수 있는 여러가지 방법이 있겠지만 저는 딕셔너리를 이용하여 해결하였습니다. 다음으로 친구가 몇명이 있는지 알아야 합니다. 유니온 파인드는 부모를 통해 연결 여부만 파악했었는데 추가적으로 친구가 몇명이 있는지 저장해 놓는 로직이 필요합니다. 그럼 직접 코드를 작성하면서 문제를 해결해 보도록 하겠습니다.

입력 받기

T = int(input())

for _ in range(T):
    F = int(input())
    network = {}
    for _ in range(F):
        a, b = input().split()
        print(union(a, b))

첫째 줄에 테스트 케이스 T를 입력 받습니다. 총 T번 반복하면서 각각의 문제를 해결합니다.

먼저 친구 관계의 수 F를 입력 받습니다. 테스트 케이스마다 새로 관계를 만들어 주어야 하기 때문에 네트워크 관계를 나타내는 network 라는 딕셔너리를 만들어 초기화 합니다.

다음으로 친구의 이름 a, b를 입력 받습니다. 숫자가 아니기 때문에 map을 사용할 필요가 없습니다. 그리고 두 친구를 union 해서 친구의 수를 출력하도록 합니다.

딕셔너리 초기화 하기

문제를 확인해보면 친구의 목록이 처음에 주어지지 않습니다. 즉 친구 이름이 들어올 때마다 기존에 있던 친구인지, 새로운 친구인지 확인하는 부분이 필요 합니다.

def make_network(a):
    if a not in network:
        network[a] = [a, 1]

a라는 이름의 친구가 들어왔을 때 network라는 딕셔너리에 a라는 키값이 포함되어 있는 친구인지 확인해서 없으면 새로 만들어 주는 코드 입니다. a라는 친구를 만들 때 2개를 저장합니다. 앞의 a는 부모를 저장한 것입니다. 처음에는 자기 자신만 친구 네트워크에 있기 때문에 자기 자신인 a를 저장한 것입니다.

다음으로 1은 친구 네트워크안에 있는 친구의 수를 나타냅니다. 지금은 자기 자신밖에 없기 때문에 1명이라 1이라 한 것입니다.

유니온 함수

def union(a, b):
    make_network(a)
    make_network(b)

    pa, ca = find(network[a][0])
    pb, cb = find(network[b][0])
    if pa != pb:
        network[pa] = network[pb] = [pb, ca + cb]

    return network[pa][1]

기존의 유니온 함수와 비슷합니다. 다른 점은 make_network를 통해 친구가 들어왔을 때 기존에 있던 친구인지, 아닌지 체크해서 새로운 친구라면 딕셔너리에 추가해 줍니다.

다음으로 find 함수로 친구 네트워크의 부모를 찾습니다. 이 때 두 개의 출력을 받는데 pa는 find 함수로 찾은 부모 이름을, ca는 해당 네트워크의 친구의 수를 나타냅니다. pa, ca는 각각 parent_a, count_a를 뜻합니다.

a의 부모 pa와 b의 부모 pb를 비교해서 둘이 다르다면 같게 만들어 주어야 합니다. 이 때 친구들의 수 역시 저장을 합니다. 그리고 친구의 수가 저장되어 있는 network[pa][1]을 리턴합니다.

find 함수

def find(a):
    if a == network[a][0]:
        return network[a]
    else:
        network[a] = find(network[a][0])
        return network[a]

파인드 함수 역시 기존과 비슷합니다. 다만 다른점은 find 함수로 리턴하는 값은 네트워크상의 부모와 친구의 수를 리턴한다는 것입니다. 따라서 a와 network[a][0]을 비교한 것입니다. 그것만 잘 기억하면 쉽게 풀 수 있습니다.

전체 코드

전체 코드를 확인해 보겠습니다.

T = int(input())

def find(a):
    if a == network[a][0]:
        return network[a]
    else:
        network[a] = find(network[a][0])
        return network[a]

def make_network(a):
    if a not in network:
        network[a] = [a, 1]

def union(a, b):
    make_network(a)
    make_network(b)

    pa, ca = find(network[a][0])
    pb, cb = find(network[b][0])
    if pa != pb:
        network[pa] = network[pb] = [pb, ca + cb]

    return network[pa][1]

for _ in range(T):
    F = int(input())
    network = {}
    for _ in range(F):
        a, b = input().split()
        print(union(a, b))