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

예전에 C++로 풀다가, 테스트케이스만 다 맞추고, 채점결과에서 정확성 문제 실패해서 포기했던 문제다. 요번에 파이썬으로 다시 풀어봐야겠다!

문제 설명

n개의 송전탑이 전선을 통해 하나의 트리 형태로 연결되어 있습니다. 당신은 이 전선들 중 하나를 끊어서 현재의 전력망 네트워크를 2개로 분할하려고 합니다. 이때, 두 전력망이 갖게 되는 송전탑의 개수를 최대한 비슷하게 맞추고자 합니다.

송전탑의 개수 n, 그리고 전선 정보 wires가 매개변수로 주어집니다. 전선들 중 하나를 끊어서 송전탑 개수가 가능한 비슷하도록 두 전력망으로 나누었을 때, 두 전력망이 가지고 있는 송전탑 개수의 차이(절대값)를 return 하도록 solution 함수를 완성해주세요.

제한 사항

• n은 2 이상 100 이하인 자연수입니다.
• wires는 길이가 n-1인 정수형 2차원 배열입니다.
  • wires의 각 원소는 [v1, v2] 2개의 자연수로 이루어져 있으며, 이는 전력망의 v1번 송전탑과 v2번 송전탑이 전선으로 연결되어 있다는 것을 의미합니다.
  • 1 ≤ v1 < v2 ≤ n 입니다.
  • 전력망 네트워크가 하나의 트리 형태가 아닌 경우는 입력으로 주어지지 않습니다.

입출력 예

n	wires	result
9	[[1,3],[2,3],[3,4],[4,5],[4,6],[4,7],[7,8],[7,9]]	3
4	[[1,2],[2,3],[3,4]]	0
7	[[1,2],[2,7],[3,7],[3,4],[4,5],[6,7]]	1

입출력 예 설명

입출력 예 #1

• 다음 그림은 주어진 입력을 해결하는 방법 중 하나를 나타낸 것입니다.

• 

• 4번과 7번을 연결하는 전선을 끊으면 두 전력망은 각 6개와 3개의 송전탑을 가지며, 이보다 더 비슷한 개수로 전력망을 나눌 수 없습니다.
• 또 다른 방법으로는 3번과 4번을 연결하는 전선을 끊어도 최선의 정답을 도출할 수 있습니다.

입출력 예 #2

• 다음 그림은 주어진 입력을 해결하는 방법을 나타낸 것입니다.

• 

• 2번과 3번을 연결하는 전선을 끊으면 두 전력망이 모두 2개의 송전탑을 가지게 되며, 이 방법이 최선입니다.

입출력 예 #3

• 다음 그림은 주어진 입력을 해결하는 방법을 나타낸 것입니다.

• 

• 3번과 7번을 연결하는 전선을 끊으면 두 전력망이 각각 4개와 3개의 송전탑을 가지게 되며, 이 방법이 최선입니다.

풀이

• 예전 C++ 풀이를 보면, 입력 데이터 변환하는 것부터 노가다였다.

#include <vector>
#include <algorithm>

using namespace std;

int solution(int n, vector<vector<int>> wires) {
    int answer = 9999999;

    vector<int> graph[n+1];
    vector<char> visited(n+1, ' ');

    sort(wires.begin(), wires.end());
    
    for(int i=0;i<wires.size();i++){
        graph[wires[i][0]].push_back(wires[i][1]);
        graph[wires[i][1]].push_back(wires[i][0]);
    }
    
    for(int i=1;i<=n;i++){
       sort(graph[i].begin(), graph[i].end());
    }
    
    for(int i=0;i<wires.size();i++){
        fill(visited.begin(), visited.end(), ' ');
        visited[wires[i][0]] = 'L';
        visited[wires[i][1]] = 'R';

        int left, right, sum = 0;

        while(sum < n){
            for(int z=1;z<=n;z++){
                for(int y=0;y<graph[z].size();y++){
                    if(visited[graph[z][y]] == ' ') visited[graph[z][y]] = visited[z];
                }
            }
            for(int z=n;z>0;z--){
                for(int y=0;y<graph[z].size();y++){
                    if(visited[graph[z][y]] == ' ') visited[graph[z][y]] = visited[z];
                }
            }
            
            left = 0;
            right = 0;
            for(int z=1;z<=n;z++){
                if(visited[z]=='L')left++;
                if(visited[z]=='R')right++;
            }
            sum = left + right;
            answer = min(answer, abs(left - right));
        }
    }

    return answer;
}

• 테스트 케이스는 다 맞췄는데...

테스트 1
입력값 〉	9, [[1, 3], [2, 3], [3, 4], [4, 5], [4, 6], [4, 7], [7, 8], [7, 9]]
기댓값 〉	3
실행 결과 〉	테스트를 통과하였습니다.

테스트 2
입력값 〉	4, [[1, 2], [2, 3], [3, 4]]
기댓값 〉	0
실행 결과 〉	테스트를 통과하였습니다.

테스트 3
입력값 〉	7, [[1, 2], [2, 7], [3, 7], [3, 4], [4, 5], [6, 7]]
기댓값 〉	1
실행 결과 〉	테스트를 통과하였습니다.

테스트 4
입력값 〉	8, [[1, 2], [1, 3], [1, 4], [4, 5], [5, 6], [6, 7], [6, 8]]
기댓값 〉	0
실행 결과 〉	테스트를 통과하였습니다.

• 채점에서 실패한다.

정확성  테스트
테스트 1 〉	통과 (0.08ms, 4.27MB)
테스트 2 〉	통과 (2.64ms, 4.31MB)
테스트 3 〉	실패 (0.12ms, 4.15MB)
테스트 4 〉	실패 (0.13ms, 4.14MB)
테스트 5 〉	실패 (0.14ms, 4.33MB)
테스트 6 〉	통과 (0.01ms, 3.66MB)
테스트 7 〉	통과 (0.01ms, 4.26MB)
테스트 8 〉	실패 (0.04ms, 3.69MB)
테스트 9 〉	통과 (0.03ms, 4.34MB)
테스트 10 〉	실패 (0.75ms, 4.31MB)
테스트 11 〉	실패 (0.72ms, 4.25MB)
테스트 12 〉	실패 (0.63ms, 4.28MB)
테스트 13 〉	실패 (0.72ms, 4.13MB)

• 근데 그 때나 지금이나 알고리즘 문제를 푸는 방식은 크게 달라진 건 없는 것 같다. (무지성ㅋㅋㅋ)
• 그리고 이거 왜 못 풀었는데 그냥 넘어갔는지 기억이 안나네...
• 코딩테스트 고득점 Kit에 있는 완전탐색 문제 7개 중에 유일하게 못 풀고 넘어갔던 문제였다.

• 문제부터 다시 제대로 이해해보기로...
• 이건 진짜 뭔가 꼼수 없이 하나씩 잘라서 해봐야하는 건가?
• 일단 시키는대로 완전탐색으로 풀어봄.

def solution(n, wires):
    answer = 100
    for i in range(n-1):
        print(i, wires[i],"을 끊어봅니다 ==================")
        heap = wires.copy()
        heap.pop(i)
        L = set()
        R = set()
        L.add(wires[i][0])
        R.add(wires[i][1])
        while heap:
            a, b = heap.pop(0)
            if a in L or b in L:
                L.add(a)
                L.add(b)
            elif a in R or b in R:
                R.add(a)
                R.add(b)
            else:
                heap.append([a, b])
        print("L:", L)
        print("R:", R)
        answer = min(abs(len(L)-len(R)), answer)
    return answer

• 테스트 케이스 모두 잘 풀림

테스트 1
입력값 〉	9, [[1, 3], [2, 3], [3, 4], [4, 5], [4, 6], [4, 7], [7, 8], [7, 9]]
기댓값 〉	3
실행 결과 〉	테스트를 통과하였습니다.
출력 〉	0 [1, 3] 을 끊어봅니다 ==================
L: {1}
R: {2, 3, 4, 5, 6, 7, 8, 9}
1 [2, 3] 을 끊어봅니다 ==================
L: {2}
R: {1, 3, 4, 5, 6, 7, 8, 9}
2 [3, 4] 을 끊어봅니다 ==================
L: {1, 2, 3}
R: {4, 5, 6, 7, 8, 9}
3 [4, 5] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 6, 7, 8, 9}
R: {5}
4 [4, 6] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5, 7, 8, 9}
R: {6}
5 [4, 7] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5, 6}
R: {8, 9, 7}
6 [7, 8] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5, 6, 7, 9}
R: {8}
7 [7, 9] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5, 6, 7, 8}
R: {9}

테스트 2
입력값 〉	4, [[1, 2], [2, 3], [3, 4]]
기댓값 〉	0
실행 결과 〉	테스트를 통과하였습니다.
출력 〉	0 [1, 2] 을 끊어봅니다 ==================
L: {1}
R: {2, 3, 4}
1 [2, 3] 을 끊어봅니다 ==================
L: {1, 2}
R: {3, 4}
2 [3, 4] 을 끊어봅니다 ==================
L: {1, 2, 3}
R: {4}

테스트 3
입력값 〉	7, [[1, 2], [2, 7], [3, 7], [3, 4], [4, 5], [6, 7]]
기댓값 〉	1
실행 결과 〉	테스트를 통과하였습니다.
출력 〉	0 [1, 2] 을 끊어봅니다 ==================
L: {1}
R: {2, 3, 4, 5, 6, 7}
1 [2, 7] 을 끊어봅니다 ==================
L: {1, 2}
R: {3, 4, 5, 6, 7}
2 [3, 7] 을 끊어봅니다 ==================
L: {3, 4, 5}
R: {1, 2, 6, 7}
3 [3, 4] 을 끊어봅니다 ==================
L: {1, 2, 3, 6, 7}
R: {4, 5}
4 [4, 5] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 6, 7}
R: {5}
5 [6, 7] 을 끊어봅니다 ==================
L: {6}
R: {1, 2, 3, 4, 5, 7}

테스트 4
입력값 〉	8, [[1, 2], [1, 3], [1, 4], [4, 5], [5, 6], [6, 7], [6, 8]]
기댓값 〉	0
실행 결과 〉	테스트를 통과하였습니다.
출력 〉	0 [1, 2] 을 끊어봅니다 ==================
L: {1, 3, 4, 5, 6, 7, 8}
R: {2}
1 [1, 3] 을 끊어봅니다 ==================
L: {1, 2, 4, 5, 6, 7, 8}
R: {3}
2 [1, 4] 을 끊어봅니다 ==================
L: {1, 2, 3}
R: {4, 5, 6, 7, 8}
3 [4, 5] 을 끊어봅니다 ==================
L: {1, 2, 3, 4}
R: {8, 5, 6, 7}
4 [5, 6] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5}
R: {8, 6, 7}
5 [6, 7] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5, 6, 8}
R: {7}
6 [6, 8] 을 끊어봅니다 ==================
L: {1, 2, 3, 4, 5, 6, 7}
R: {8}

• 최적화 해서 제출해본다.
• 성공! 진짜 완전 탐색 문제였구나...;;;

def solution(n, wires):
    answer = 100
    for i in range(n-1):
        heap = wires.copy()
        heap.pop(i)
        L, R = set(), set()
        L.add(wires[i][0])
        R.add(wires[i][1])
        while heap:
            a, b = heap.pop(0)
            if a in L or b in L:
                L.add(a)
                L.add(b)
            elif a in R or b in R:
                R.add(a)
                R.add(b)
            else:
                heap.append([a, b])
        answer = min(abs(len(L)-len(R)), answer)
    return answer
    
정확성  테스트
테스트 1 〉	통과 (2.16ms, 10.3MB)
테스트 2 〉	통과 (35.27ms, 10.1MB)
테스트 3 〉	통과 (16.33ms, 10.2MB)
테스트 4 〉	통과 (4.78ms, 10.2MB)
테스트 5 〉	통과 (10.80ms, 10.4MB)
테스트 6 〉	통과 (0.01ms, 10.2MB)
테스트 7 〉	통과 (0.01ms, 10.2MB)
테스트 8 〉	통과 (0.34ms, 10.3MB)
테스트 9 〉	통과 (0.18ms, 10.2MB)
테스트 10 〉	통과 (8.34ms, 10.2MB)
테스트 11 〉	통과 (8.90ms, 10.2MB)
테스트 12 〉	통과 (9.38ms, 10.2MB)
테스트 13 〉	통과 (8.16ms, 10.2MB)

• 고수의 풀이... 인줄 알았는데 느리다.

def solution(n, wires):
    ans = n
    for sub in (wires[i+1:] + wires[:i] for i in range(len(wires))):
        s = set(sub[0])
        [s.update(v) for _ in sub for v in sub if set(v) & s]
        ans = min(ans, abs(2 * len(s) - n))
    return ans
    
정확성  테스트
테스트 1 〉	통과 (312.04ms, 10.2MB)
테스트 2 〉	통과 (269.78ms, 10.4MB)
테스트 3 〉	통과 (221.52ms, 10.2MB)
테스트 4 〉	통과 (220.82ms, 10.3MB)
테스트 5 〉	통과 (272.53ms, 10.3MB)
테스트 6 〉	통과 (0.04ms, 10.3MB)
테스트 7 〉	통과 (0.02ms, 10.1MB)
테스트 8 〉	통과 (1.85ms, 10.1MB)
테스트 9 〉	통과 (2.72ms, 10.2MB)
테스트 10 〉	통과 (212.52ms, 10.1MB)
테스트 11 〉	통과 (282.40ms, 10.2MB)
테스트 12 〉	통과 (224.79ms, 10.2MB)
테스트 13 〉	통과 (210.24ms, 10.2MB)

• 기본적으로 속도를 조금 먹지만, 그래프가 복잡해졌을 때 빠르다.
• def find(a)란걸 만들어서 탐색하고 합친다.

uf = []

def find(a):
    global uf
    if uf[a] < 0: return a
    uf[a] = find(uf[a])
    return uf[a]

def merge(a, b):
    global uf
    pa = find(a)
    pb = find(b)
    if pa == pb: return
    uf[pa] += uf[pb]
    uf[pb] = pa

def solution(n, wires):
    global uf
    answer = int(1e9)
    k = len(wires)
    for i in range(k):
        uf = [-1 for _ in range(n+1)]
        tmp = [wires[x] for x in range(k) if x != i]
        for a, b in tmp: merge(a, b)
        v = [x for x in uf[1:] if x < 0]
        answer = min(answer, abs(v[0]-v[1]))

    return answer
    
정확성  테스트
테스트 1 〉	통과 (7.58ms, 10.2MB)
테스트 2 〉	통과 (8.56ms, 10.3MB)
테스트 3 〉	통과 (8.71ms, 10.3MB)
테스트 4 〉	통과 (4.35ms, 10.2MB)
테스트 5 〉	통과 (5.30ms, 10.2MB)
테스트 6 〉	통과 (0.03ms, 10.2MB)
테스트 7 〉	통과 (0.02ms, 10.4MB)
테스트 8 〉	통과 (0.28ms, 10.2MB)
테스트 9 〉	통과 (0.21ms, 10.4MB)
테스트 10 〉	통과 (5.35ms, 10MB)
테스트 11 〉	통과 (6.33ms, 10.2MB)
테스트 12 〉	통과 (5.73ms, 10.4MB)
테스트 13 〉	통과 (5.07ms, 10.3MB)

• 찾았다. 진짜 고수네... 이게 제일 빠르네?
• 이거 뭐지? 왤케 빨라???
• 내거 디큐로 바꿔도 이 방식보다 빠르진 않을 것 같다. pop, push를 안쓰는 방식으로 했어야...
• 암튼 코드가 길어서 긴가민가 했는데 개빠름.

def solution(n, wires):
    tree = [[] for _ in range(n + 1)]

    for a, b in wires:
        tree[a].append(b)
        tree[b].append(a)

    visited = [False] * (n + 1)
    child = [0] * (n + 1)

    """각자의 자식 개수 구하기"""
    def dfs(node):
        visited[node] = True

        for next in tree[node]:
            if not visited[next]:
                child[node] += dfs(next) + child[next]

        return 1

    dfs(1)
    result = n

    for c in child:
        result = min(result, abs(n - 2 * (c + 1)))

    return result
    
정확성  테스트
테스트 1 〉	통과 (0.12ms, 10.2MB)
테스트 2 〉	통과 (0.07ms, 10.4MB)
테스트 3 〉	통과 (0.07ms, 10.1MB)
테스트 4 〉	통과 (0.11ms, 10.3MB)
테스트 5 〉	통과 (0.07ms, 10.2MB)
테스트 6 〉	통과 (0.01ms, 10.2MB)
테스트 7 〉	통과 (0.01ms, 10.3MB)
테스트 8 〉	통과 (0.02ms, 10.3MB)
테스트 9 〉	통과 (0.02ms, 10.3MB)
테스트 10 〉	통과 (0.07ms, 10.3MB)
테스트 11 〉	통과 (0.07ms, 10.1MB)
테스트 12 〉	통과 (0.07ms, 10.1MB)
테스트 13 〉	통과 (0.07ms, 10.4MB)

• 그리고 for next는 뭐지? 이런 구문도 있구나... 그냥 변수였다. ㅡ.ㅡ;;;;;
  •  속았넹ㅋ ㅋㅋ ㅋ ㅋ
• 반복문 정리 좀 해놓자 (for, while)
  • break : 반복문 종료
  • continue : 다음 명령어를 넘기고, 이후 반복으로 (반복문 유지)
  • pass
• iter, next 함수
  • iter는 반복을 끝낼 값을 지정하면 특정 값이 나올 때 반복을 끝냅니다.
  • next는 기본값을 지정할 수 있습니다.