[백준 Python Swift] 4195번 친구 네트워크

4195번 친구 네트워크

• 문제바로가기

풀이 방법

• Union-Find 알고리즘을 사용한다
• find : 어떤 원소가 주어졌을 때 이 원소의 부모(대표)를 반환한다
• union : 두 원소의 부모를 동일하게 만든다(두 집합을 하나의 집합으로 합친다)
• find_parent 함수를 하나 만들어 재귀적으로 최종 부모를 찾는다
• union 함수를 하나 만들어 오른쪽의 최종 부모를 왼쪽의 최종 부모로 만든다
  • 그리고 오른쪽의 친구 수를 왼쪽의 친구 수에 더해준다
• 마지막으로 왼쪽의 친구 수를 출력하면 된다

풀이

Python

def find_parent(x):
    if parent[x] == x:
        return x
    else:
        last_parent = find_parent(parent[x])
        parent[x] = last_parent
        return parent[x]

def union(x, y):
    parent_x = find_parent(x)
    parent_y = find_parent(y)
    if parent_x != parent_y:
        parent[parent_y] = parent_x
        count[parent_x] += count[parent_y]
    
test_case = int(input())
for _ in range(test_case):
    f = int(input())
    parent = {}
    count = {}
    for _ in range(f):
        left, right = input().split(" ")
        if left not in parent:
            parent[left] = left
            count[left] = 1
        if right not in parent:
            parent[right] = right
            count[right] = 1
        union(left, right)
        print(count[find_parent(left)])

Swift

let testCase = Int(readLine()!)!

for _ in 1...testCase {
    let f = Int(readLine()!)!
    var parent: [String: String] = [:]
    var count: [String: Int] = [:]
    
    func findParent(_ x: String) -> String {
        if parent[x] == x {
            return x
        } else {
            let lastParent = findParent(parent[x]!)
            parent[x] = lastParent
            return parent[x]!
        }
    }
    func union(_ x: String, _ y: String) -> Void {
        let parentX = findParent(x)
        let parentY = findParent(y)
        
        if parentX != parentY {
            parent[parentY] = parentX
            count[parentX]! += count[parentY]!
        }
    }
    
    for _ in 1...f {
        let names = readLine()!.split(separator: " ")
        let left = String(names[0]), right = String(names[1])
        
        if parent.keys.contains(left) == false {
            parent[left] = left
            count[left] = 1
        }
        if parent.keys.contains(right) == false {
            parent[right] = right
            count[right] = 1
        }
        
        union(left, right)
        print(count[findParent(left)]!)
    }
}