파이썬으로 풀어보는 백준 1949번: 우수 마을

https://www.acmicpc.net/problem/1949

1949번: 우수 마을

 

 

내 풀이:

def DFS(current):
    visited[current] = True
    s = [current]
    DP[current][0] = population[current]
    while s:
        current = s[-1]
        for next in connection[current]:
            if not visited[next]:
                s.append(next)
                visited[next] = True
                DP[next][0] = population[next]
                break
 
        else:
            s.pop()
            if s:
                DP[s[-1]][0] += DP[current][1]
                DP[s[-1]][1] += max(DP[current][0], DP[current][1])
 
 
N = int(input())
population = [0] + list(map(int,input().split()))
connection = [[] for _ in range(N+1)]
for i in range(N-1):
    a, b = map(int,input().split())
    connection[a].append(b)
    connection[b].append(a)
 
DP = [[0, 0] for _ in range(N+1)]
visited = [False] * (N+1)
DFS(1)
print(max(DP[1][0], DP[1][1]))

Python3, 500ms

문제가 제법 까다로웠다. 트리 형태의 문제지만 그리디 알고리즘을 이용하여 문제를 풀어야 했다. 문제 조건 중 '우수 마을'로 선정되지 못한 마을은 적어도 하나의 '우수 마을'과는 인접해 있어야 한다는 조건이 있는데 '우수 마을'로 선정된 마을 주민 수의 총합을 최대로 해야 한다는 조건을 만족시키기 위해서 자연스럽게 해결되는 부분이기 때문에 세 번째 조건은 무시해줘도 된다. 문제를 풀 때 sys를 import 하는 것은 지양하고 있기 때문에 처음에는 재귀로 구현했다가 파이썬 재귀 제한(1000번)에 걸려 런타임 에러가 났다. 이를 해결하기 위해 어쩔 수 없이 DFS를 스텍 형태로 구현하였는데 for-else문을 이용하는 이와 같은 방식은 stack이 많이 생길수록 시간이 느려지는 현상이 발생한다. 때문에 백준 기준으로 python을 이용해 문제를 푼 사람들 중 속도가 매우 느린 편으로 제출되었지만 모든 사람들이 sys를 이용하여 재귀 제한을 늘려줬기 때문에 스택 형태의 구현도 연습해볼 필요가 있다고 생각한다. 특히 트리에서 부모를 찾기 위한 아이디어로 스택을 pop() 해준 순간 s[-1]이 pop을 해준 현재의 부모가 된다는 점일 이용해준 것이 주요했다.