[코드트리] 노드의 값 C++

문제

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

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

 

코드

#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

const int MAX_N = 10001;

// 주어지는 정보
int n;
int weight[MAX_N]; // weight[i]: i번 노드에 적혀있는 번호
vector<int> edges[MAX_N];

// 풀이를 위해 선언한 정보
int parent[MAX_N]; // i번 노드의 부모노드
int prefix_weight[MAX_N]; // prefix_weight[i]: i노드를 루트로 하는 서브트리의 가중치 누적합
int num_of_node[MAX_N]; // num_of_node[i]: i번 노드를 루트로 하는 서브트리의 노드 개수
bool visited[MAX_N]={false, };

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

    prefix_weight[x]=weight[x];
    num_of_node[x]=1;

    for(int i=0; i<(int)edges[x].size(); ++i){
        int y = edges[x][i];
        // y의 부모가 x라면 dp[x]에 dp[x]+dp[y] 저장
        if(parent[y]==x){
            prefix_weight[x]+=prefix_weight[y];
            num_of_node[x]+=num_of_node[y];
        }
    }
}

int main() {
    // input
    // 정점의 수 n
    cin >> n;
    // 1번부터 n번까지 n 개의 정점에 적혀 있는 정수가 번호 순서대로 공백으로 구분되어 주어집니다.
    for(int i=1; i<=n; ++i){
        cin >> weight[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);

    // output
    int answer=0;
    for(int i=1; i<=n; ++i){
        answer+= abs(num_of_node[i]-prefix_weight[i]);
    }
    cout << answer << '\n';
    return 0;
}

처음에 작성했던 코드

#include <iostream>
#include <vector>
using namespace std;

const int MAX_N = 10001;

int n, max_node;
// 각 정점에는 음이 아닌 정수가 적혀 있고, n 개의 정점에 적혀있는 모든 정수의 합은 정확하게 n
int node_num[MAX_N]; // node_num[i]: i번 노드에 적혀있는 번호
vector<int> edges[MAX_N];
int parent[MAX_N]; // i번 노드의 부모노드
int dp[MAX_N]; // dp[i]: i를 루트로 하는 서브트리에서 "노드에 적힌 수가 0인 노드"의 갯수
bool visited[MAX_N]={false, };

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

    dp[x] = (node_num[x]==0) ? 1 : 0;

    for(int i=0; i<(int)edges[x].size(); ++i){
        int y = edges[x][i];
        // y의 부모가 x라면 dp[x]에 dp[x]+dp[y] 저장
        if(parent[y]==x){
            dp[x]+= dp[y];
        }
    }
}

int main() {
    // input
    // 정점의 수 n
    cin >> n;
    // 1번부터 n번까지 n 개의 정점에 적혀 있는 정수가 번호 순서대로 공백으로 구분되어 주어집니다.
    for(int i=1; i<=n; ++i){
        cin >> node_num[i];
        if(node_num[i] > node_num[max_node]) max_node=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[max_node]=true;
    DFS(max_node);

    // output
    int answer=0;
    for(int i=1; i<=n; ++i){
        if(node_num[i]==0){
            answer+= dp[i];
        }
    }
    cout << answer << '\n';
    return 0;
}