17831번 - 대기업 승범이네(python3)

문제 링크: https://www.acmicpc.net/problem/17831

17831번: 대기업 승범이네

또 하나의 트리DP 문제입니다. 어째 DP는 잘 못 푸는데 트리DP에 꽂힌 것 같아요... 

 

Tree 구조가 주어져 있고, 각 node마다 어떤 양의 정수값을 가지고 있습니다. 두 노드를 그룹으로 만들고, 두 노드가 가지고 있는 값을 곱해서 그룹의 값으로 정의할건데, 가능한 그룹의 값 총합의 최댓값을 구하는 문제입니다. 이 때 그룹에 속하는 두 노드 중 parent를 멘토, child를 멘티라고 합시다. 각 노드는 두 개 이상의 그룹에 동시에 들어갈 수 없고, 한 그룹에 속하거나 아니면 속한 그룹이 없거나 2가지 경우가 가능합니다.

 

우선, 다음은 문제 푸는데 사용한 코드입니다.

 

2533번 - 사회망 서비스(SNS) 문제에서 각 노드가 얼리어답터인지, 아닌지 2가지 경우가 가능했던 것처럼, 이문제에서는 각 노드가 그룹에 속하느냐, 속하지 않느냐 2가지 경우가 가능합니다. 각 노드마다 2가지 경우에 대한 값을 모두 구해줘야 합니다. Node v가 그룹에 속할 때 v를 root로 하는 subtree의 그룹값 총 합의 최댓값을 dp_mat[v][0], 그룹에 속하지 않을 경우 최댓값을 dp_mat[v][1]에 저장해둡니다.

 

현재 노드 v가 있고, v의 child 집합 child={a, b, c}가 있다고 합시다. 편의 상 v의 임의의 child를 x라고 쓰겠습니다. 노드 v가 멘토가 아닐 경우에는 a, b, c가 멘토든 아니든 상관 없으므로 dp_mat[x][0]과 dp_mat[x][1] 중 큰 값을 선택해서 더한 후 dp_mat[v][1]에 assign하면 됩니다.

 

만약 노드 v가 멘토일 경우(dp_mat[v][0])에는 a,b,c 중 하나를 멘티로 선정하게 됩니다. 만약 a를 멘티로 선정하게 된다면, dp_mat[b], 및 dp_mat[c]에서는 가장 큰 값을 취사선택하고, dp_mat[a]에서는 무조건 dp_mat[a][1]을 선택해야 합니다. a는 멘티면서 멘토일 수 없으니까요. 그렇다면 dp_mat[v][0]를 구하기 위해서는 모든 x에 대해 완성된 dp_mat[x]를 알고 있어야 합니다.

 

그런데 tree에서 재귀를 돌릴 때, line 21에서 for문을 돌면서 nei 내 순서대로 dp_mat[x]가 완성되게 됩니다. 만약 a, b, c 순서대로 함수 dp가 수행된다면, nei==a인 경우에 아직 dp_mat[b], dp_mat[c]는 완성된 상태가 아닙니다. a를 멘티로 삼는 경우에 dp_mat[x][0]값을 바로 구할 수 없다는 것이죠. 그러면 line 21번의 for문을 모두 돌고 나서(모든 x에 대해 dp_mat[x]가 완성된 후) 다시 for문을 돌아야 할까요? 그렇게 해 봤는데 메모리 초과가 뜨더라구요...

 

이런 문제를 해결하기 위해 con이라는 변수를 사용했습니다. Line 21번의 for문을 돌면서 v가 선택되지 않는 경우인 dp_mat[v][1]가 완성되고, 이 값은 dp_mat[x]에서 큰 값을 취사선택해서 더한 값입니다. 만약 a를 멘티로 선정하게 된다면?

dp_mat[v][1]에서 먼저 stat[v-1]*stat[a-1](v-a 그룹의 값)을 더하고,
    1. dp_mat[a][0]>dp_mat[a][1]인 경우에는(a가 멘토인 경우가 optimal이었을 경우) dp_mat[a][0]을 빼고 dp_mat[a][1]을 더해주면 되고((v-a)그룹 + a가 멘토가 아닐 때 최댓값),
    2. dp_mat[a][0]<dp_mat[a][1]인 경우에는 더 이상 아무런 작업을 할 필요가 없습니다. 

그러니까, 각 child x에 대해 1번의 경우 stat[v-1]*stat[a-1] - dp_mat[a][0] + dp_mat[a][1]을, 2번의 경우 stat[v-1]*stat[a-1]을 구하고 이들 중 가장 optimal한 값을 주는 x를 찾으면 되는겁니다. 위 2가지 경우를 line 36처럼 con = stat[v-1]*stat[a-1] - "max(dp_mat[a][0]-dp_mat[a][1], 0)" 1줄로 만들 수 있습니다. 가장 큰 con값을 주는 x를 멘티로 삼는 경우를 dp_mat[v][0]로 assign하면 됩니다. Line 21번 for문을 돌면서 구한 dp_mat[v][1]에 optimal con을 더해주면 됩니다.

 

설명이 좀 장황한데, 그림으로 그려보면 이해가 좀 더 쉽습니다. 시간이 많으면 ppt로 그림이라도 그려보고 싶은데.... 시간이 부족해서... 코드에 주석 자세히 달았으니 코드 참고해주세요:)