백준[17831] 대기업 승범이네

문제 설명

1. 문제 링크

http://icpc.me/17831

17831번: 대기업 승범이네

2. 풀이

트리에서 dp를 통해 풀 수 있다.

dp[i][0] : 노드 i가 멘티가 아닐 때, 최댓값

dp[i][1] : 노드 i가 멘티일 때, 최댓값이라 하자.

dp[i][0] = max(dp[i의 자식 노드][0], dp[i의 자신노드][1] + cost[i] * cost[i의 자식노드)의 총합 (이때, 멘티는 최대 1명만 선택 가능)

dp[i][1] = dp[i의 자식 노드][0]의 총 합

멘티를 한명 이하로만 선택하는 구현을 하기 위해 멘티가 i번째를 멘티로 할 때, (멘티가 하나도 없는 모든 자식 합) - (i가 멘티가 아닐 경우) + (i가 멘티일 경우)를 확인하며 구현한다.

3. 코드

#include <iostream>
#include <cstring>
#include <vector>

using namespace std;

const int N = 2e5 + 5;
int cost[N];
int dp[N][2];
vector<int> v[N];

int dpf(int cur, int menti) {
    int &ret = dp[cur][menti];
    if (~ret) return ret;
    ret = 0;
    int noMenti = 0;
    if (!menti) {
        for (auto nxt : v[cur]) noMenti += dpf(nxt, 0);
        ret = noMenti;
        for (auto nxt : v[cur]) {
            ret = max(ret, noMenti - dpf(nxt, 0) + dpf(nxt, 1) + cost[cur] * cost[nxt]);
        }
    } else {
        for (auto nxt : v[cur]) ret += dpf(nxt, 0);
    }
    return ret;
}

int main() {
    int n; cin >> n;
    for (int i = 2; i <= n; i++) {
        int ipt; cin >> ipt;
        v[ipt].push_back(i);
    }
    for (int i = 1; i <= n; i++) cin >> cost[i];
    memset(dp, -1, sizeof dp);
    cout << dpf(1, 0);
}C++