[코드트리] 인접한 노드2 C++

문제

https://www.codetree.ai/missions/9/problems/adjacent-node-2?&utm_source=clipboard&utm_medium=text

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

코드

저번에 풀었던 문제랑 비슷한 거 같아서 1번의 DFS로 해결하려고 해서 잘 안 풀렸다. 아직 재귀는 쉽지 않은 거 같다...

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int MAX_N = 10001;

int n;
vector<int> edges[MAX_N];
bool visited[MAX_N]={false, };
int num[MAX_N]; // num[i]: i번 노드에 적힌 값

int dp[MAX_N][2]; // dp[i][j]: j=0 선택x 경우, j=1 선택o 경우
vector<int> picked;

void FillDpTable(int x){
    dp[x][1]=num[x];

    for(int i=0; i<(int)edges[x].size(); ++i){
        int y =edges[x][i];
        if(visited[y])continue;

        visited[y]=true;
        FillDpTable(y);

        // x번 노드가 포함된 경우 y번 노드는 포함될 수 없습니다.
        // x번 노드가 포함되지 않은 경우 y번 노드는 포함될 수도,
        // 포함되지 않을 수도 있습니다.
        dp[x][1] += dp[y][0];
        dp[x][0] += max(dp[y][0], dp[y][1]);
    }
}

void PickNode(int x, int parent, bool pickFlag){
    // 노드 선택 여부 결정
    if (!pickFlag && dp[x][1] > dp[x][0]) {
        picked.push_back(x);
        pickFlag = true;
    }
    else pickFlag = false; 

    // 자식 노드 탐색
    for(int i=0; i<(int)edges[x].size(); ++i){
        int y=edges[x][i];
        if(y==parent) continue;
        PickNode(y, x, pickFlag);
    }
}

int main() {
    //input
    cin >> n;
    for(int i=1; i<=n; ++i){
        cin >> num[i];
    }
    for(int i=1; i<n; ++i){
        int a, b;
        cin >> a >> b;
        // 간선 연결
        edges[a].push_back(b);
        edges[b].push_back(a);
    }

    // solution
    visited[1]=true;
    FillDpTable(1);

    PickNode(1, 0, false);

    sort(picked.begin(), picked.end());

    // output
    cout << max(dp[1][0], dp[1][1]) << '\n';
    for(auto e: picked){
        cout << e << " ";
    }
    return 0;
}

 

처음 코드

반례는 인접한 정점에서 썼던 예제가 있다.

5
5 3 4 10 2
1 2
1 3
2 4
3 5

 

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int MAX_N = 10001;

int n;
vector<int> edges[MAX_N];
bool visited[MAX_N]={false, };
int num[MAX_N]; // num[i]: i번 노드에 적힌 값

int dp[MAX_N][2]; // dp[i][j]: j=0 선택x 경우, j=1 선택o 경우
vector<int> picked;

void DFS(int x){
    dp[x][1]=num[x];

    for(int i=0; i<(int)edges[x].size(); ++i){
        int y =edges[x][i];
        if(visited[y])continue;

        visited[y]=true;
        DFS(y);

        // x번 노드가 포함된 경우 y번 노드는 포함될 수 없습니다.
        // x번 노드가 포함되지 않은 경우 y번 노드는 포함될 수도,
        // 포함되지 않을 수도 있습니다.
        dp[x][1] += dp[y][0];
        dp[x][0] += max(dp[y][0], dp[y][1]);
    }

    // x를 선택하는 게 나은 선택이라면 x 저장
    if(dp[x][1]>dp[x][0]) {
        picked.push_back(x);
    }
}

int main() {
    //input
    cin >> n;
    for(int i=1; i<=n; ++i){
        cin >> num[i];
    }
    for(int i=1; i<n; ++i){
        int a, b;
        cin >> a >> b;
        // 간선 연결
        edges[a].push_back(b);
        edges[b].push_back(a);
    }

    // solution
    visited[1]=true;
    DFS(1);
    sort(picked.begin(), picked.end());

    // output
    cout << max(dp[1][0], dp[1][1]) << '\n';
    for(auto e: picked){
        cout << e << " ";
    }
    return 0;
}