[프로그래머스] 코딩테스트 연습 - 모두 0으로 만들기 (Python)

코딩테스트 연습 - 모두 0으로 만들기 [월간 코드 챌린지 시즌2]

코딩테스트 연습 - 모두 0으로 만들기

Stack를 이용한 DFS를 사용해서 문제를 해결하기가 까다로운 문제였습니다..

 

풀이

먼저 문제의 해결방법은 "LeafNode부터 ParentNode의 값과 합하며 Root로 이동"입니다.

해당 방법이 가능한 이유는 아래와 같습니다.

1. 노드의 가중치를 변경하는 연산은 Tree의 전체 합에 영향을 주지 않습니다.
2. 모든 가중치를 0으로 만드는 +,- 연산들은 교환 법칙이 성립합니다.

위의 이유들로 인해서 모든 가중치의 합이 0이 아니라면 모든 가중치를 0으로 만들 수 없고, LeafNode들부터 ParentNode에 값을 더하며 Root로 이동한다면 최소 연산수를 구할 수 있는 것이죠.

 

그러나 DFS를 사용한다면 몇 가지 신경 써야 할 부분들이 있습니다.

• Stack을 이용한 DFS

Stack을 이용한 DFS를 사용한다면 Root부터 시작하여 Leaf를 찾고 Leaf노드에서 Root까지 다시 올라와야 하는 작업이 필요하기 때문에 저는 해당 방법으로 시도했다가 3개의 테스트 케이스에서 시간 초과가 나서 방법을 조금 바꾸었습니다.

Node의 차수가 1이라면 해당 Node는 LeafNode이므로 먼저 LeafNode들을 찾아 Queue에 넣고 ParentNode로 가중치를 합하며 이동하도록 수정해서 해결할 수 있었습니다.

 

• Recursive를 이용한 DFS

Node의 개수가 최대 300,000개 이므로 Python의 Default Recursive Depth로는 해결하지 못할 테스트 케이스가 있습니다.

그래서 Recursive Depth를 늘려주기 위해 아래의 코드를 사용하면 문제를 해결할 수 있습니다.

import sys
sys.setrecursionlimit(300000)

 

 

Code - Queue

from collections import defaultdict, deque
def solution(a, edges):
    if sum(a) != 0: return -1
    answer = 0
    mdict = defaultdict(set)
    for s, e in edges:              # Tree의 형태 저장
        mdict[s].add(e)
        mdict[e].add(s)

    queue, visit = deque(), set()
    for k, v in mdict.items():      # LeafNode 찾기
        if len(v)==1:
            queue.append(k)

    while queue:                    # Leaf부터 Parent로 이동
        curr = queue.popleft()
        visit.add(curr)
        for node in mdict[curr]:
            if node not in visit:
                a[node] += a[curr]
                answer += abs(a[curr])
                mdict[node].remove(curr)    # Parent와 자신의 연결 삭제
                if len(mdict[node])==1:     # Parent가 Leaf가 되었을 때
                    queue.append(node)
    return answer

 

 

Code - Recursive

from collections import defaultdict
def solution(a, edges):
    if sum(a) != 0: return -1
    mdict = defaultdict(list)
    for s, e in edges:              # Tree의 형태 저장
        mdict[s].append(e)
        mdict[e].append(s)
        
    def dfs(curr=0, prev=0, answer=0):
        for node in mdict[curr]:
            if node != prev:        # Child로 재귀
                answer += dfs(node, curr)
        a[prev] += a[curr]
        answer += abs(a[curr])
        return answer
    
    return dfs()

 

 

 

문제 설명

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.

• 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다.

하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다. 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요. (만약 처음부터 트리의 모든 정점의 가중치가 0이라면, 0을 return 해야 합니다.)

 

 

제한사항

• a의 길이는 2 이상 300,000 이하입니다.
  • a의 모든 수는 각각 -1,000,000 이상 1,000,000 이하입니다.
  • a[i]는 i번 정점의 가중치를 의미합니다.
• edges의 행의 개수는 (a의 길이 - 1)입니다.
  • edges의 각 행은 [u, v] 2개의 정수로 이루어져 있으며, 이는 u번 정점과 v번 정점이 간선으로 연결되어 있음을 의미합니다.
  • edges가 나타내는 그래프는 항상 트리로 주어집니다.

 

 

입출력 예

a	edges	result
[-5,0,2,1,2]	[[0,1],[3,4],[2,3],[0,3]]	9
[0,1,0]	[[0,1],[1,2]]	-1

 

 

입출력 예 #1

• 다음 그림은 주어진 트리의 모든 정점의 가중치를 0으로 만드는 과정을 나타낸 것입니다.

1. 2번 정점과 3번 정점을 선택하여 2번 정점은 1 감소시키고, 3번 정점은 1 증가시킵니다. (2번 반복)
2. 3번 정점과 4번 정점을 선택하여 4번 정점은 1 감소시키고, 3번 정점은 1 증가시킵니다. (2번 반복)
3. 0번 정점과 3번 정점을 선택하여 3번 정점은 1 감소시키고, 0번 정점은 1 증가시킵니다. (5번 반복)

• 모든 정점의 가중치를 0으로 만드는 데 필요한 최소 행동 횟수는 9번이므로, 9를 return 해야 합니다.

입출력 예 #2

• 주어진 트리는 모든 정점의 가중치를 0으로 만드는 것이 불가능하므로, -1을 return 해야 합니다.