만득

코딩 공부 정리용

​

© 2021. All rights reserved.

백준 4195 친구 네트워크

23 Jun 2021

https://www.acmicpc.net/problem/4195

친구 관계가 생긴 순서대로 주어졌을 때 두사람의 친구 네트워크에 몇 명이 속해있는지 구하는 문제이다. 두 친구가 속한 집합이 하나로 합쳐졌을 때 집합 전체의 크기를 구하는 문제라 할 수 있다.

집합간의 관계를 구하는 문제이기 때문에 union find를 사용할 예정이다. 이 때 각 집합은 번호로 매겨져있는 것이. 아닌 string으로 들어오게 된다. 따라서 dictionary를 이용하여 string을 int 번호로 바꿔주는 작업이 필요하다. 새로운 이름이 나올 때마다 번호를 부여하고 부모로 본인 자신을 가르키는 집합을 추가해주면 된다.

추가로 union할 때 집합 전체의 크기를 알려줘야 하기 때문에 parent 정보를 저장하는 배열 외에 집합의 크기를 저장해두는 배열을 하나 더 사용해야한다.

struct DisjointSet {
    var nameToInt:[String:Int] = [:]
    var arr:[Int] = []
    var memberCount:[Int] = []
    func nodeNumber(_ str: String) -> Int {
        guard let tmp = nameToInt[str] else {
            return -1
        }
        return tmp
    }
    mutating func insert(_ str: String) -> Int {
        let count = arr.count
        arr.append(count)
        memberCount.append(1)
        nameToInt[str] = count
        return count
    }
    mutating func findRoot(_ n: Int) -> Int {
        if arr[n] == n {
            return n
        }
        else {
            let root = findRoot(arr[n])
            arr[n] = root
            return root
        }
    }
    mutating func union(_ a: Int, _ b: Int) -> Int {
        let rootA = findRoot(a), rootB = findRoot(b)
        if rootA < rootB {
            arr[rootB] = rootA
            memberCount[rootA] += memberCount[rootB]
            return memberCount[rootA]
        }
        else if rootA > rootB{
            arr[rootA] = rootB
            memberCount[rootB] += memberCount[rootA]
            return memberCount[rootB]
        }
        else {
            return memberCount[rootA]
        }
    }
}

let testcase = Int(readLine()!)!
var str = ""
for _ in 1...testcase {
    let n = Int(readLine()!)!
    var disjointSet = DisjointSet()
    for _ in 1...n {
        let relation = readLine()!.split(separator: " ").map{String($0)}
        let person1 = relation[0], person2 = relation[1]
        var number1 = disjointSet.nodeNumber(person1)
        var number2 = disjointSet.nodeNumber(person2)
        if number1 == -1 {
            number1 = disjointSet.insert(person1)
        }
        if number2 == -1 {
            number2 = disjointSet.insert(person2)
        }
        str += "\(disjointSet.union(number1, number2))\n"
    }
}
print(str)

Related posts