[프로그래머스] 모두 0으로 만들기, 월간 코드 챌린지, python

programmers.co.kr/learn/courses/30/lessons/76503

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

대회 당일에는 트리로 만들면 될거같긴한데 그이후에 어떻게 접근해야할지 몰랐다.

if sum(a): return -1 만 해놓고 그이후엔 트리로 풀면되지 않을까하고 뚱땅뚱땅하긴했지만 33점인가 나왔다.

---

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

def solution(a, edges):
    # sum이 0이 아니면 가중치바꿔도 남을듯?
    if sum(a):
        return -1
    # greedy
    # 부모 자식간의 관계를 트리로 생성
    cp_link = [list() for i in range(len(a))]
    # cp_link[0] 은 부모노드, 최상위 노드혹은 없는 경우 -1
    visited = [0] * len(a)
    for i in range(len(edges)):
        cp_link[edges[i][0]].append(edges[i][1])
        cp_link[edges[i][1]].append(edges[i][0])

    global answer
    answer = 0

    # 후위 연산(중위인가?)으로 leaf node부터 연산
    def preorder(node):
        global answer
        # 무향 그래프를 그대로 사용하므로 방문한 노드인지 체크
        if visited[node]:
            return 0
        # 방문하지 않았으면 방문 체크
        visited[node] = 1
        # 인접한 모든 노드 방문
        for i in range(len(cp_link[node])):
            a[node] += preorder(cp_link[node][i])
        # 모든 노드 방문했으면 위로 올라가
        tmp = a[node]
        a[node] = 0
        answer += abs(tmp)
        return tmp

    preorder(0)
    if not a[0]:
        return answer
    return -1

이문제를 오늘 다시 풀었고 트리를 후위연산인지 중위 연산인지 좀 애매하게 사용하긴 했지만 예제는 잘 풀렸다.

제출했더니 런타임 에러... 

대충 예상을 해보면 최악의 경우 최대 범위의 값이 들어가면 answer의 크기가 int형으로는 부족해서 메모리 초과가 뜬거같다.

 

---

def solution(a, edges):
    # sum이 0이 아니면 가중치바꿔도 남을듯?
    if sum(a):
        return -1
    # greedy
    # 부모 자식간의 관계를 트리로 생성
    cp_link = [list() for i in range(len(a))]
    # cp_link[0] 은 부모노드, 최상위 노드혹은 없는 경우 -1
    visited = [0] * len(a)
    for i in range(len(edges)):
        cp_link[edges[i][0]].append(edges[i][1])
        cp_link[edges[i][1]].append(edges[i][0])
    global answer
    answer = 0
    stack = []
    stack.append(0)
    visited[0] = 1
    son = 0
    while stack:
        node = stack[-1]
        if son != 0:
            a[node] += son
            son = 0
        for i in range(len(cp_link[node])):
            if visited[cp_link[node][i]]: continue
            visited[cp_link[node][i]] = 1
            stack.append(cp_link[node][i])
            break
        else: # leaf node인 경우 혹은 인접노드를 다 확인한 경우
            son = a[node]
            a[node] = 0
            answer += abs(son)
            stack.pop()

    if not a[0]:
        return answer
    return -1

그래서 순회를 재귀말고 반복문으로 풀었는데 다른 테스트 케이스에 시간초과가 뜬다. 

---

테스트케이스만 보면 답은 다 맞는것으로 나오니 알고리즘은 맞는다는 뜻이고 python에서는 자료형의 크기때문에 틀린게 아니라는 것이다.

혹시나해서 재귀 깊이를 늘렸더니 바로 성공했다.

import sys
sys.setrecursionlimit(10**7)
def solution(a, edges):
    if sum(a):
        return -1
    # greedy
    # 부모 자식간의 관계를 트리로 생성
    cp_link = [list() for i in range(len(a))]
    # cp_link[0] 은 부모노드, 최상위 노드혹은 없는 경우 -1
    visited = [0] * len(a)
    for i in range(len(edges)):
        cp_link[edges[i][0]].append(edges[i][1])
        cp_link[edges[i][1]].append(edges[i][0])
    global answer
    answer = 0

    # 트리 순회. leaf node부터 연산
    def preorder(node):
        global answer
        # 무향 그래프를 그대로 사용하므로 방문한 노드인지 체크
        if visited[node]:
            return 0
        # 방문하지 않았으면 방문 체크
        visited[node] = 1
        # 인접한 모든 노드 방문
        for i in range(len(cp_link[node])):
            a[node] += preorder(cp_link[node][i])
        # 모든 노드 방문했으면 위로 올라가
        son = a[node]
        a[node] = 0
        answer += abs(son)
        return son

    preorder(0)
    if not a[0]:
        return answer
    return -1