[백준 1949번] 우수 마을 (Python3)

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)

def DFS(last,now):
  DP1[now] += population[now]
  diff = []
  for next in graph[now]:
    if next == last:
      continue
    DFS(now,next)
    DP23next = max(DP2[next],DP3[next])
    DP1[now] += DP23next
    DP2[now] += max(DP1[next],DP3[next])
    DP3[now] += DP23next
    diff.append(DP1[next]-DP23next)
  if diff:
    diff.sort(reverse=True)
    if diff[0] < 0:
      DP3[now] += diff[0]
    else:
      for d in diff:
        if d < 0:
          break
        DP3[now] += d

N = int(input())
population = [*map(int,input().split())]

graph = [[] for i in range(N)]
for i in range(N-1):
  x,y = map(int,input().split())
  graph[x-1].append(y-1)
  graph[y-1].append(x-1)

DP1,DP2,DP3 = [0]*N,[0]*N,[0]*N    #우수,비우수,비우수+옆집우수

DFS(0,0)
print(max(DP1[0],DP2[0],DP3[0]))

백준 2533번 사회망 서비스(SNS)와 거의 유사한 문제이다. 그런데 차이점이 있다면 우수마을에 대한 조건이 "인접해선 안된다" "하나의 우수마을과는 인접해야한다" 로 두개가 존재한다는 점이다. 그래서 어떻게 해결할까 고민하다가 DP를 3개로 만드는 방식으로 해결하였다.

 

알고리즘을 요약하면,

1. DP1은 now가 우수마을일 때의 최대값, DP2는 비우수마을일때 최대값, DP3는 비우수마을이면서 최소 하나의 우수마을과 접할때의 최대값을 의미한다.

2. DP1[now]는 DP2[next]와 DP3[next] 간의 최댓값

3. DP2[now]는 DP1[next]와 DP3[next] 간의 최댓값 (DP2[next]는 우수마을과 한번도 접하지 않았으므로 now가 우수마을이어야한다.)

4. DP3[now]는 최소 하나이상 접하는 최댓값,

5. 4에서 DP3를 구하는 방식은 DP2[next],DP3[next] 중 최댓값을 일단 더한 후, DP1[next]의 최댓값과의 차이를 저장한다.

6. 최댓값과의 차이가 양수인 것이 존재하면 양수인걸 전부 더해주고, 전부 음수면 절댓값이 가장 작은 값만 더해준다.

7. DFS 역추적으로 최댓값을 구하고 출력한다.

 

그런데 사실 이 문제는 굳이 DP를 3개나 만들 필요가 없다고 한다. 왜냐하면 조건 1,2,3 중 1,2만 만족하면 3번은 자명하게 만족하기 때문이다.

이걸 알고 문제를 풀면 2533번과 마찬가지로 DP 2개만으로 문제를 풀 수가 있다.

 

 

오늘의 교훈) 문제의 조건 중 자명하게 만족하는 조건을 제외할 수 있어야 한다.