[Swift] BOJ 2606 바이러스

BOJ 2606 바이러스

 

접근 방식

 

1번 컴퓨터와 연결된 컴퓨터 수를 세는 문제. (solved.ac 기준 실버 3)

 

인접 리스트로 그래프를 만들고 1번에서 DFS로 도달 가능한 노드를 셈. 시작 노드 1은 감염 대상이 아니므로 방문 수에서 하나를 빼고 출력함(a.count - 1).

 

주의할 점은 입력을 받는 부분. 그래프는 간선이 양방향이라 서로가 서로를 인접 리스트에 담아야 함. [Int: [Int]] 형태로 받을 때, 2번과 5번이 서로 연결돼 있다면

 

• 2: [1,3,5]
• 5: [1,2,6]

 

처럼 양쪽 모두에 넣어야 함. 트리처럼 한 방향만 담다가 종종 실수함.

 

 

나는 DFS로 풀었는데 BFS를 사용해도 상관없음.

 

코드

 

import Foundation

struct b2606 {
    static func run() {
        // Input 처리
        var n = Int(readLine()!)!
        var input = Int(readLine()!)!
        
        var list = [Int: [Int]]()
        
        for _ in 0..<input {
            let line = readLine()!.split(separator: " ").map{ Int(String($0))! }
            let start = line[0]
            let end = line[1]
            
            if list[start] == nil {
                list[start] = [end]
            } else {
                list[start]?.append(end)
            }
            
            if list[end] == nil {
                list[end] = [start]
            } else {
                list[end]?.append(start)
            }
        }
        
//        print(list)
        let a = DFS(graph: list, start: 1)
        print(a.count-1)
    }
}

import Foundation

func DFS<T> (graph: [T: [T]], start: T) -> [T] {
    var visitedQueue: [T] = []
    var needVisitStack: [T] = [start]
    
    while !needVisitStack.isEmpty {
        let node: T = needVisitStack.removeLast()
        if visitedQueue.contains(node) { continue }
        
        visitedQueue.append(node)
        needVisitStack += graph[node] ?? []
    }
    
    return visitedQueue
}