[백준 17831번] 파이썬 - 대기업 승범이네

백준 17831 - 대기업 승범이네

시간 제한 1초, 메모리 제한 512MB

# 조건

• ㈜승범이네는 사장 승범이를 포함한 N명의 직원이 모두 판매원인 다단계 회사이다.
• 사장 승범이를 제외한 모든 판매원에게는 사수가 한 명씩 배정된다.
• 만약 판매원 A가 B의 사수라면, B를 A의 부사수라고 부른다.
• 작년에 창설된 ㈜승범이네는 큰 수익률을 기록하고 대기업으로 거듭났다.
• ㈜승범이네의 더 큰 성장을 위해 멘토링 제도를 도입하려고 한다.
• 승범이는 사수와 부사수 관계에 있는 두 판매원을 각각 서로의 멘토와 멘티로 만들 수 있으며, 이 경우 두 판매원이 멘토링 관계에 있다고 한다.
• 한 판매원은 최대 1개의 멘토링 관계에만 속할 수 있다.
  • 즉, 한 판매원이 여러 명의 멘토가 되거나, 여러 명의 멘티가 되거나, 멘토인 동시에 멘티가 될 수는 없다.
  • 물론 멘토링 관계에 속하지 않는 직원이 있을 수도 있다.
• 이렇게 만들어진 멘토링 관계에서는 시너지 효과가 발생한다. 승범이는 모든 판매원의 실력을 수치화시켰으며, 한 멘토링 관계에서 발생하는 시너지는 멘토와 멘티의 실력의 곱과 같다는 것을 발견했다.
• 승범이는 적절하게 멘토링 관계를 만들어, 모든 멘토링 관계에서 발생하는 시너지의 합을 최대로 만들려고 한다.

• 위 그림은 ㈜승범이네의 회사 구조와 멘토링 관계를 나타낸 예시이다.
• 각 원은 한 명의 판매원을 의미하며, 원 안에 쓰인 숫자는 그 판매원의 실력이다.
• 화살표는 사수-부사수 관계를 나타낸다. 이 경우 가능한 시너지의 최대 합은 5×7 + 4×3 + 3×3 + 4×5 + 3×1 = 79이다.

입력

• 첫 번째 줄에 판매원들의 수 N(2 ≤ N ≤ 200,000)이 주어진다.
• 판매원들은 1번, 2번, …, N번으로 번호가 매겨지며, 승범이는 항상 1번이다.
• 두 번째 줄에 2번 판매원부터 N번 판매원의 사수가 순서대로 공백으로 구분되어 주어진다.
• 세 번째 줄에 i번 판매원의 실력을 나타내는 정수 A1, A2, …, AN (0 ≤ Ai ≤ 100)이 순서대로 공백으로 구분되어 주어진다.

출력

• 첫 번째 줄에 모든 멘토링 관계에서 발생하는 시너지의 합의 최댓값을 출력한다.

# 접근 방법

• dp를 이용하여 풀어주면 된다.
• 이 문제의 핵심은 현재 노드가 리프 노드가 아닌 경우
  • 나를 포함한 서브 트리에서의 최댓값과
  • 나를 제외한 서브 트리에서의 최댓값을 기록해주며
  • 갱신해 나가면 된다.
• 그렇기 때문에 dp는 [[0,0] for _ in range(N)]의 크기로 생성해주고
  • 각 인덱스의 0번은 나를 포함하지 않은 경우
  • 1번은 나를 포함한 경우로 기록해준다.
• 따라서, dp 값을 갱신할 때는
  • 0번 인덱스의 경우 자식 노드의 max(dp[i][0], dp[i][1]) 을 더해주며 갱신하면 되고
  • 1번 인덱스의 경우 max(dp[i][1], dp[i][0] - max(dp[j]) + dp[j][0] + skill[i] * skill[j])의 점화식을 구할 수 있다.
    • 즉, 위에서 기록한 나를 포함하지 않는 경우에서 자식 노드의 max값을 다시 빼주고
    • 자식 노드가 포함되지 않은 자식 노드의 0번 인덱스와
    • 나와 자식 노드를 곱한 시너지 값을 현재 dp[i][1]과 비교해주면 된다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
sys.setrecursionlimit(10**6)  

def check(node):  
    # 리프 노드의 부모노드까지 들어가주기  
    # 나를 포함하지 않은 시너지는 자식 노드의 max값 더해주기   
    for i in tree[node]:  
        check(i)  
        dp[node][0] += max(dp[i])  

    # 나를 포함한 값 갱신  
    # 나를 포함하지 않은 값 - 현재 자식 노드의 max값 + 내 자식 노드를 포함하지 않은 최댓값 + 나와 자식 노드의 시너지    
    for j in tree[node]:  
        dp[node][1] = max(dp[node][1], dp[node][0] - max(dp[j]) + dp[j][0] + skill[node] * skill[j])  

N = int(input())  
tree = [[] for _ in range(N)]  
parent = [*map(int, input().split())]  
for i, j in enumerate(parent):  
    tree[j-1].append(i+1)  

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

# 0번 -> 나를 포함하지 않은 시너지  
# 1번 -> 나를 포함한 시너지  
dp = [[0, 0] for _ in range(N)]  
check(0)  
print(max(dp[0]))