[고득점 Kit(탐욕법)] 섬 연결하기 (Swift, Python)

코딩테스트 연습 - 섬 연결하기

 

풀이 

 섬을 연결할 때 일단 가장 적은 비용의 다리를 먼저 사용해야 하므로 입력으로 주어진 costs를 비용을 기준으로 오름차순 정렬을 한다. 그 후 차례대로 다리를 사용해가면서 이미 연결된 두 섬에 대한 다리인지를 체크한다. 이미 두 섬이 연결되어 있으면 해당 다리를 사용하지 않고 넘어간다. 

 다리를 사용할때마다 몇 개의 다리를 사용했는지 체크한다. n개의 섬이 모두 연결되기 위한 최소 다리 개수는 n-1개이기 때문에 다리가 n-1개가 사용됐다면 탐색을 종료한다.

 

 

코드

[Swift]

import Foundation

var edges: [[Int]] = Array(repeating: [Int](), count: 100) // n번 섬과 직접적으로 연결된 섬들
var visited: [Bool] = Array(repeating: false, count: 100) // dfs 탐색할 때 사용하기 위한 방문 체크 배열
var flag: Bool = false // dfs에 사용하기 위한 flag 변수

// x섬과 y섬이 이미 연결되어있는지 체크하기 위한 dfs 함수
func dfs(_ x: Int, _ y: Int) {
    // dfs를 돌면서 이미 x섬을 방문했거나 처음 주어진 x섬과 y섬이 연결되어 있는지 체크를 완료했다면 탐색 종료
    if visited[x] || flag {
        return
    }
    visited[x] = true
    
    // 현재 x섬에 연결된 섬을 탐색하면서 y섬을 만나면 flag를 변경하고 종료
    // 그렇지 않을 경우 계속해서 dfs 호출
    for i in edges[x] {
        if i == y {
            flag = true
            break
        }
        
        dfs(i, y)
    }
}

func solution(_ n:Int, _ costs:[[Int]]) -> Int {
    var answer: Int = 0
    
    // x섬과 y섬이 연결되어 있는지를 기록하기 위한 2차원 배열
    // true일 경우 굳이 dfs를 돌리지 않는다.
    var link: [[Bool]] = Array(repeating: Array(repeating: false, count: n), count: n)
    
    let sortedCosts: [[Int]] = costs.sorted { $0[2] < $1[2] } // cost를 기준으로 오름차순 정렬
    var edgeCount: Int = 0 // 현재 섬을 연결하는데 사용된 다리의 개수
    
    for i in sortedCosts {
        let x: Int = i[0]
        let y: Int = i[1]
        let cost: Int = i[2]
        
        if link[x][y] { // x섬과 y섬이 연결되어 있다면 continue
            continue
        }
        
        // dfs 탐색을 위한 visited 배열 false 초기화
        for j in 0..<n {
            visited[j] = false
        }
        
        dfs(x, y) // x섬과 y섬 연결 체크
        
        // dfs 탐색 결과 연결되어 있다면 탐색 종료
        if flag {
            flag = false
            continue
        }
        
        edges[x].append(y)
        edges[y].append(x)
        link[x][y] = true
        link[y][x] = true
        
        edgeCount += 1
        answer += cost
        
        // 다리가 n-1개가 된 순간 탐색 종료
        if edgeCount == n-1 {
            break
        }
    }
    
    
    return answer
}

[Python]

edges = [[] for _ in range(100)] # n번 섬과 직접적으로 연결된 섬들
visited = [False] * 100 # dfs 탐색할 때 사용하기 위한 방문 체크 배열 
flag = False # dfs에 사용하기 위한 flag 변수

# x섬과 y섬이 이미 연결되어있는지 체크하기 위한 dfs 함수 
def dfs(x, y):
    global flag

    # dfs를 돌면서 이미 x섬을 방문했거나 처음 주어진 x섬과 y섬이 연결되어 있는지 체크를 완료했다면 탐색 종료
    if visited[x] == True or flag:
        return
    visited[x] = True

    # 현재 x섬에 연결된 섬을 탐색하면서 y섬을 만나면 flag를 변경하고 종료 
    # 그렇지 않을 경우 계속해서 dfs 호출 
    for i in edges[x]:
        if i == y:
            flag = True
            break

        dfs(i, y)


def solution(n, costs):
    answer = 0

    # x섬과 y섬이 연결되어 있는지를 기록하기 위한 2차원 배열
    # true일 경우 굳이 dfs를 돌리지 않는다.
    link = [[False] * n for _ in range(n)]

    costs.sort(key=lambda x: x[2]) # cost를 기준으로 오름차순 정렬
    edgeCount = 0 # 현재 섬을 연결하는데 사용된 다리의 개수

    for (x, y, cost) in costs:
        if link[x][y]: # x섬과 y섬이 연결되어 있다면 continue
            continue
        
        # dfs 탐색을 위한 visited 배열 false 초기화
        for i in range(n):
            visited[i] = False

        dfs(x, y) # x섬과 y섬 연결 체크

        global flag
        if flag: # dfs 탐색 결과 연결되어 있다면 탐색 종료 
            flag = False
            continue

        edges[x].append(y)
        edges[y].append(x)
        link[x][y] = link[y][x] = True

        edgeCount += 1
        answer += cost

        # 다리가 n-1개가 된 순간 탐색 종료 
        if edgeCount == n-1:
            break

    return answer