[프로그래머스] 섬 연결하기(Python)

https://programmers.co.kr/learn/courses/30/lessons/42861

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

 

두세 시간 걸려서 풀었던 문제다. 알고리즘을 어떻게 짜야할지 오래 고민했다

 

내가 짠 알고리즘을 아래 기준으로 설명하겠다.

cost = [[0,1,1],[0,2,2],[1,2,5],[1,3,1],[2,3,8]]

n = 4로 주어졌을 때

일단 각 섬의 정보를 모아놓는 리스트 안에 python set을 통해 각 set에 현재 연결된 섬의 정보를 담는다.

island = [set() for _ in range(n)]

 다음으로 섬 연결 비용 기준으로 정렬한다.

costs.sort(key = lambda x:x[2])

이렇게 하면 costs는 [[0, 1, 1], [2, 4, 1], [2, 3, 1], [3, 4, 1], [0, 4, 5]] 와 같이 정렬된다.

 

for cost in costs:
        check = -1
        for i in range(n):
            if check == -1:
                if len(island[i]) == 0:
                    island[i].update([cost[0],cost[1]])
                    answer += cost[2]
                    check = i
                    break
                elif cost[0] in island[i] and cost[1] in island[i]:
                    break
                elif cost[0] in island[i] or cost[1] in island[i]:
                    island[i].update([cost[0],cost[1]])
                    answer += cost[2]
                    check = i
            else :
                for j in range(n):
                    if len(island[j]) ==0 :
                        break
                    elif cost[0] in island[j] or cost[1] in island[j]:
                        island[check].update(island[j])
        if len(island[check]) == n:
            return answer

 

여기서부터 costs를 for loop를 돌며 cost [0] , cost [1]의 섬 정보를 가지고

미리 선언 해 놓은 island 리스트를 순회하며 다리 정보가 업데이트됐는지 확인한다.

 

여기서

1. 순회하고 있는 set의 길이가 0 일 때 (현재 set에 섬이 아무것도 없을 때)

    현재 순회하는 set에 섬 번호를 추가하고 cost 비용을 더한다.

2. 순회하고 있는 set에 두 섬이 모두 등록되어 있다면 (이미 연결된 다리 이므로 추가 X)

3. 순회하고 있는 set에 두 섬 중 하나만 있다면

    현재 set에 두 섬을 업데이트시켜준다.

    check 값을 현재 set 번호를 등록시켜줘서 혹시 다른 set에 이번에 추가한 섬 정보가 있다면 두 set의 섬 정보를 합쳐준다

 

전체 코드

def solution(n, costs):
    island = [set() for _ in range(n)]
    costs.sort(key = lambda x:x[2])
    answer = 0
    for cost in costs:
        check = -1
        for i in range(n):
            if check == -1:
                if len(island[i]) == 0:
                    island[i].update([cost[0],cost[1]])
                    answer += cost[2]
                    check = i
                    break
                elif cost[0] in island[i] and cost[1] in island[i]:
                    break
                elif cost[0] in island[i] or cost[1] in island[i]:
                    island[i].update([cost[0],cost[1]])
                    answer += cost[2]
                    check = i
            else :
                for j in range(n):
                    if len(island[j]) ==0 :
                        break
                    elif cost[0] in island[j] or cost[1] in island[j]:
                        island[check].update(island[j])
        if len(island[check]) == n:
            return answer
    return answer