대기업 승범이네

• 트리 DP 문제
• 각 서브트리에 대해서, 가질수 있는 최댓값 (dp[x])와, 루트가 child와 멘토링을 안맺었을때의 최댓값(dp2[x])를 구해놓으면 된다.
• 그러면 어떤 트리에 대해서 dp와 dp2값을 서브트리들에 대한 값으로부터 계산하는 것은 아래처럼 된다.
  • $dp2[root] = \sum_i{dp[i]}$
  • $dp[root] = max(dp2[root], max_k(\sum_{i≠k}{dp[i]} + dp2[k] + A[root]*A[k]))$
• 후자는 그냥 저렇게 단순히 계산하면 O(|child|^2)일것 같지만, 아래처럼 바꾸면 O(|child|)이 된다
  • $dp[root] = \sum_i{dp[i]} + max(0, max_k(-dp[k] + dp2[k] + A[root]*A[k]))$
• 전체 시간복잡도는 O(n)

"""Solution code for "BOJ 17831. 대기업 승범이네".

- Problem link: https://www.acmicpc.net/problem/17831
- Solution link: http://www.teferi.net/ps/problems/boj/17831

Tags: [DP] [DP on Tree]
"""

from teflib import ttree


def main():

    def calc_dp(subtree_vals, root):
        val_without_root = max_inc_with_root = 0
        for child, sub_val, sub_val_without_root in subtree_vals:
            val_without_root += sub_val
            inc = A[root] * A[child] - sub_val + sub_val_without_root
            max_inc_with_root = max(max_inc_with_root, inc)
        return (root, val_without_root + max_inc_with_root, val_without_root)

    N = int(input())
    bosses = [int(x) for x in input().split()]
    A = [int(x) for x in input().split()]

    tree = [[] for _ in range(N)]
    for i, boss in enumerate(bosses, 1):
        tree[i].append(boss - 1)
        tree[boss - 1].append(i)

    print(ttree.dp_on_tree(tree, calc_dp, 0)[1])


if __name__ == '__main__':
    main()

• Dependency: teflib.ttree.dp_on_tree