BOJ 2132 나무 위의 벌레

📝 문제 정보

링크: https://www.acmicpc.net/problem/2132

🧐 관찰 및 접근

트리에서 특정 경로 위에서 노드들의 가중치의 합
- 트리의 경로란, 하나의 간선을 두번 지나지 않으며 $a_i \in V, (a_i, a_{i+1}) \in E$ 를 만족하는 집합
- 다시말해, 그냥 특정 간선을 두번 타지 않고 갈 수 있는 시점~종점의 길
어떻게 구할지 생각해보자
- 어떤 점 $a_0$에서 경로를 시작한다고 해보자.
- 그 다음에 갈 수 있는 점은 $a_0$의 자식중 한곳에만 갈 수 있다.
  - 한 자식한테 들어간다면, 다시 나올수가 없기 때문에
- 그 자식을 $a_1$이라고 한다면, 마찬가지로 그 자식중 한곳에만 갈 수 있다.
- 어? 이거 DFS아닌가?
DFS처럼 할 수 있는 것 같다.
- 그러면 자식중에 어디로 가야하는가?
  - 물론 가장 열매를 많이 먹을 수 있는 곳으로!

💻 풀이

코드 (C++):

int N;
vector<int> links[10010];
int val[10010];

int dfs(int cur, int par){
    int ret = 0;
    for(auto nxt: links[cur]){
        if(nxt == par) continue;
        ret = max(ret, dfs(nxt, cur));
    }
    return ret + val[cur];
}

void solve(){
    cin >> N;
    rep(i, 1, N+1) cin >> val[i];
    rep(i, 0, N-1){
        int u, v; cin >> u >> v;
        links[u].push_back(v);
        links[v].push_back(u);
    }

    int ans = 0, idx = 1;
    rep(i, 1, N+1){
        int res = dfs(i, -1);
        if(res > ans){
            ans = res;
            idx = i;
        }
    }
    cout << ans << ' '<< idx << '\n';
}

코드(C++) Tree DP를 사용한 $O(N)$ 풀이:


int N;
vector<int> links[10010];
int apple[10010];

struct Node{
    /*
    이때 subtree는 두가지 경우가 있음
    1. 자식 노드의 subtree의 값이 더 큰경우
    2. 자식 노드 두개까지 경로의 가중치 합 + 해당 노드의 가중치 합이 더 큰 경우
    */
    pii subtree; // 해당 노드를 지나는 서브트리에서 최적 값 (최대 사과 개수, 최소 인덱스)
    pii leaf; // 해당 노드 ~ 리프 노드 사이에서 최적 값 (최대 사과 개수, 최소 인덱스)
    // 간단한 max 연산을 위해 인덱스는 음수로 저장
};
Node DP[10010];

void dfs(int cur, int par){
    pii bestSub = {apple[cur], -cur};
    pii bestLeaf = {apple[cur], -cur};

    vector<pii> childLeafs;
    for(auto nxt: links[cur]){
        if(nxt == par) continue;
        dfs(nxt, cur);
        bestSub = max(bestSub, DP[nxt].subtree); // 1번 경우

        // 가장 이득인놈 두개만 남기면 됨
        childLeafs.push_back(DP[nxt].leaf);
        sort(all(childLeafs), greater<pii>());
        if((int)childLeafs.size() > 2) childLeafs.pop_back();
    }

    if((int)childLeafs.size() >= 2){ // 2번 경우
        pii cand = {apple[cur] + childLeafs[0].first + childLeafs[1].first, max(childLeafs[0].second, childLeafs[1].second)};
        bestSub = max(bestSub, cand);
    }
    if((int)childLeafs.size() >= 1){
        pii cand = {apple[cur] + childLeafs[0].first, childLeafs[0].second};
        bestLeaf = max(bestLeaf, cand);
    }

    pii cand = {bestLeaf.first, max(bestLeaf.second, -cur)};
    bestSub = max(bestSub, cand);

    DP[cur].subtree = bestSub;
    DP[cur].leaf = bestLeaf;
}

void solve(){
    cin >> N;
    rep(i, 1, N+1) cin >> apple[i];
    rep(i, 0, N-1){
        int u, v; cin >> u >> v;
        links[u].push_back(v);
        links[v].push_back(u);
    }

    dfs(1, -1);
    cout << DP[1].subtree.first << ' ' << -DP[1].subtree.second << '\n';
}

끝점 처리 로직이 상당히 까다롭네요. 값만 신경써도 되면 쉬웠을텐데…