백준 16437번 양 구출 작전 C++

https://www.acmicpc.net/problem/16437

16437번: 양 구출 작전

 

각 섬에서 1번 섬으로 갈 수 있는 유일한 경로가 하나는 꼭 있고 연결되어있는 섬은 각 섬당 하나씩 주어지기 때문에 1번 섬이 루트 노드인 트리로 나타낼 수 있다.

 

각 섬들을 돌면서 1번 섬으로 갈 수 있는 양의 수를 구하면 되는데, 그냥 구하면 시간초과가 난다.

 

내가 사용한 시간초과를 해결하기 위한 방법은 두가지이다.

 

1. 1번 섬으로 가는 도중 만나는 양을 데리고 간다.

2. 늑대가 없거나 배가 불러서 더이상 먹을 수 없는 경로를 저장해서 그 경로에 진입한다면 별도의 검사 없이 바로 1번 섬으로 보낸다.

 

#include <iostream>
#include <string>
#include <vector>
#include <math.h>
#include <algorithm>
#include <utility>
#include <stack>
#include <queue>
#include <math.h>
#include <set>
#include <map>
#include <list>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <limits.h>

using namespace std;
using int64 = long long;

enum class AnimalType {
    SHEEP,
    WOLF
};

struct Node {
    AnimalType animalType;
    int64 cnt;
    int parent;
};

int n;
vector<Node> v;
vector<bool> CanSkip;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    cin >> n;
    v = vector<Node>(n + 1);
    CanSkip = vector<bool>(n + 1, false);
    for (int i = 2; i <= n; i++) {
        char type;
        int64 cnt;
        int parent;
        cin >> type >> cnt >> parent;
        AnimalType animalType = type == 'S' ? AnimalType::SHEEP : AnimalType::WOLF;
        v[i] = { animalType,cnt,parent };
    }
    int64 ans = 0;
    for (int i = 2; i <= n; i++) {
        if (v[i].animalType == AnimalType::WOLF) continue;
        if (v[i].cnt == 0) continue;
        if (CanSkip[i]) {
            ans += v[i].cnt;
            continue;
        }
        vector<int> passes;
        passes.push_back(i);
        int parent = v[i].parent;
        int64 sheepCnt = v[i].cnt;
        bool canGo = true;
        while (true) {
            if (CanSkip[parent]) break;
            passes.push_back(parent);
            if (parent == 1) break;
            if (v[parent].animalType == AnimalType::SHEEP) {
                sheepCnt += v[parent].cnt;
                v[parent].cnt = 0;
                parent = v[parent].parent;
                continue;
            }
            if (v[parent].cnt >= sheepCnt) {
                v[parent].cnt -= sheepCnt;
                sheepCnt = 0;
                canGo = false;
                break;
            }
            sheepCnt -= v[parent].cnt;
            v[parent].cnt = 0;
            parent = v[parent].parent;
        }
        ans += sheepCnt;
        v[i].cnt = 0;
        if (canGo) {
            for (auto a : passes) {
                CanSkip[a] = true;
            }
        }
    }
    cout << ans;
}