[프로그래머스][python]전력망을 둘로 나누기

문제

https://school.programmers.co.kr/learn/courses/30/lessons/86971

소스코드 

from collections import deque
import math

def solution(n, wires):
    answer = 101
    wires.sort()
    for i in range(n-1):
        tree = [[] for _ in range(n + 1)]
        for j in range(n-1): # 트리 만들기
            if i == j : # 간선 하나 자르기
                continue
            node1 = wires[j][0]
            node2 = wires[j][1]
            tree[node1].append(node2)
            tree[node2].append(node1)

        visited = [False for _ in range(n+1)] # 방문한 노드 체크 할 리스트
        res = [0 for _ in range(n+1)] # 방문한 노드 root 노드 체크

        q = deque()
        node1 = wires[i][0]
        node2 = wires[i][1]
        q.append(node1) #자른 간선을 기준으로 체크
        q.append(node2) #자른 간선을 기준으로 체크
        visited[node1] = True # 시작 노드 방문 체크
        visited[node2] = True # 시작 노드 방문 체크
        res[node1] = node1
        res[node2] = node2

        while q:
            cur_node = q.popleft()
            for next_node in tree[cur_node]:
                if not visited[next_node]:
                    visited[next_node] = True
                    res[next_node] = res[cur_node]
                    q.append(next_node)
                    
        ans1 = res.count(node1)
        ans2 = res.count(node2)
        answer = min(answer,abs(ans1-ans2))

    return answer

설명

BFS로 Union-Find를 구현해서 탐색한 경우

이중 포문을 돌려서 트리를 만들어서 통과는했지만 가장 오래 걸린 경우는 10.01ms까지 갔다 ;;; 

소스코드2

from collections import deque
import math

def solution(n, wires):
    answer = float('inf')
    tree = [[] for _ in range(n + 1)] #트리를 우선 생성
    for w1,w2 in wires:
        tree[w1].append(w2)
        tree[w2].append(w1)

    for w1,w2 in wires: #노드를 미리 체크한다 (두 노드는 이미 체크했으므로 더이상 체크 안함 -> 연결 안됨)
        visited = [False for _ in range(n + 1)]  # 방문한 노드 체크 할 리스트
        res = [0 for _ in range(n + 1)]  # 방문한 노드 root 노드 체크

        q = deque()
        node1 = w1
        node2 = w2
        q.append(node1)  # 자른 간선을 기준으로 체크
        q.append(node2)  # 자른 간선을 기준으로 체크
        visited[node1] = True  # 시작 노드 방문 체크
        visited[node2] = True  # 시작 노드 방문 체크
        res[node1] = node1
        res[node2] = node2

        while q:
            cur_node = q.popleft()
            for next_node in tree[cur_node]:
                if not visited[next_node]:
                    visited[next_node] = True
                    res[next_node] = res[cur_node]
                    q.append(next_node)

        ans1 = res.count(node1)
        ans2 = res.count(node2)
        answer = min(answer, abs(ans1 - ans2))

    return answer

설명

BFS + UnionFind 는 동일하지만

이중포문 때문에 시간이 오래 걸렸을 거라고 판단했다.

간선을 끊는다의 의미를 다시 생각해보았다.

주어진 노드를 visited 처리함으로써 끊어줌 (체크하면 다음번에 체크 안하니까)

소스코드3

def dfs(tree,visited,v):
    cnt = 1
    for next_n in tree[v]:
        if not visited[next_n]: #방문하지 않았고 연결되어있다면
            visited[next_n] = True #방문 체크하고
            cnt += dfs(tree,visited,next_n)
    return cnt

def solution(n, wires):
    answer = float('inf')
    tree = [[] for _ in range(n + 1)]
    for n1,n2 in wires:
        tree[n1].append(n2)
        tree[n2].append(n1)

    for n1,n2 in wires:
        visited = [False for _ in range(n + 1)]  # 방문한 노드 체크 할 리스트 초기화
        visited[n1] = True # 노드를 탐색했다고 체크함으로써 끊어버리기
        visited[n2] = True
        ans = dfs(tree,visited,n1) #한쪽만 체크하기 (전체에서 한쪽 체크한거 빼면 반대쪽)

        answer = min(answer, abs(ans - (n-ans))) #한섹션 - 다른섹션

    return answer

설명

dfs를 사용해서 체크해줌 시간이 많이 줄었다

훨씬 빠르다. DFS를 써서 시간이 줄었나? 라고 생각해보니까 노드 수가 적기 때문에 그럴 수가 없다 

그리고 속도는 DFS, BFS를 선택하는 기준이 아니다.

다시 생각해보니까 BFS로 탐색할 땐 양쪽 노드를 다 체크해주는 방면 DFS는 한쪽 노드만 체크했다.

그럼 BFS도 한쪽만 탐색해주면 되지 않을까? 라는 생각에 다시 코드를 짜보았다.

 

소스코드4

from collections import deque
import math

def solution(n, wires):
    answer = float('inf')
    tree = [[] for _ in range(n + 1)] #트리를 우선 생성
    for w1,w2 in wires:
        tree[w1].append(w2)
        tree[w2].append(w1)

    for w1,w2 in wires: #노드를 미리 체크한다 (두 노드는 이미 체크했으므로 더이상 체크 안함 -> 연결 안됨)
        visited = [False for _ in range(n + 1)]  # 방문한 노드 체크 할 리스트
        res = [0 for _ in range(n + 1)]  # 방문한 노드 root 노드 체크

        q = deque()
        node1 = w1
        node2 = w2
        q.append(node1)  # 자른 간선을 기준으로 체크 **한쪽만 체크해주기 위해 한쪽 루트 노드만 넣음**
        visited[node1] = True  # 시작 노드 방문 체크
        visited[node2] = True  # 시작 노드 방문 체크 **반대쪽 끊어줌**
        cnt1 = 1 #한쪽 섹션 카운트


        while q:
            cur_node = q.popleft()
            for next_node in tree[cur_node]:
                if not visited[next_node]:
                    visited[next_node] = True
                    q.append(next_node)
                    cnt1 +=1

        answer = min(answer, abs(cnt1 - (n-cnt1))) # 체크한쪽 - 반대쪽 (굳이 안구해도됨)

    return answer

설명 

1. 이중 포문으로 돌리지 않으면 시간이 줄어든다.

2. 주어진 간선을 visited 처리한다 = 탐색하지 않으니 간선을 끊어준다는 뜻

3. 양쪽을 다 체크하지 말고 한쪽만 체크해서 전체에서 빼주면 된다.

4.bfs 사용

 

이 3가지를 모두 충족시켜서 bfs를 사용해서 알고리즘을 적용하니 가장 최적의 알고리즘을 짰다!

후기

4번의 리팩토링을 통해 시간을 많이 줄였다.

10.01ms (소스1)-> 2.73ms(소스2) -> 0.38ms(소스3) -> 0.08ms(소스4)