260507 TIL - 양과 늑대 문제와 이진트리 백트래킹 최적화

2026. 5. 7. 10:36언리얼 7기 본캠프

어제 심화반에서 풀었던 '양과 늑대' 문제에서, 풀었던 방법과 문제점, 극한 케이스의 경우에 대해 알아보려 함

 

https://school.programmers.co.kr/learn/courses/30/lessons/92343

 

프로그래머스

SW개발자를 위한 평가, 교육의 Total Solution을 제공하는 개발자 성장을 위한 베이스캠프

programmers.co.kr

 

그리고 아래 코드가 수업 도중에 풀었던 방식의 코드임

 

#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>

using namespace std;

void backtracking(const vector<int>& info, unordered_map<int, vector<int>>& tree, int& answer, int wolf, int sheep, int node, vector<int> candidates){
    //양, 늑대 갱신 후 늑대 수 검증
    int cursheep = info[node] == 0 ? sheep + 1 : sheep;
    int curwolf = info[node] == 1 ? wolf + 1 : wolf;
    if (curwolf >= cursheep) return;
    
    //정답 갱신(현재 양과 기존 최대 양 중 큰 값)
    answer = max(cursheep, answer);
    
    //목적지 후보에서 스스로를 삭제 후 자신에서 출발하는 목적지들을 목적지에 추가
    candidates.erase(remove(candidates.begin(), candidates.end(), node), candidates.end());
    for(int i : tree[node]) candidates.push_back(i);
    
    //목적지를 돌면서 재귀. 목적지를 &가 아닌 복사본으로 넘기며 자연스럽게 함수가 끝나면 소멸되게(백트래킹) 함
    for(int i : candidates)
    {
        backtracking(info, tree, answer, curwolf, cursheep, i, candidates);
    }
}

int solution(vector<int> info, vector<vector<int>> edges) {
    int answer = 0;
    
    //이진트리 연결 정보를 저장
    unordered_map<int, vector<int>> tree;
    for(const vector<int>& v : edges) tree[v[0]].push_back(v[1]);
    
    backtracking(info, tree, answer, 0, 0, 0, {});
    
    return answer;
}

 

문제 내 이미지를 통해 이 방식의 해설을 덧붙이자면

 

파란색 화살표 방향대로 진행시 예제

 

이 문제의 경우 '노드의 이동제약이 없다'는 것이 포인트로, 첨부한 이미지처럼 1번에 간 이후 무조건 2번이나 4번으로 가야하는것이 아닌, 8번으로도 이동할 수 있는것이 특징임. 기존 DFS같은 방식을 사용할 경우 여기서 1번에 갈 경우 2번 또는 4번으로 가는것이 불가능하기 때문에, 계속해서 후보가 담긴 vector를 복사해가며 그 후보들을 돌면서 탐색할경우 최대 양을 구할 수 있는 결과값에 도달할 수 있음. 늑대의 수가 양의 수보다 같거나 많아지면 즉시 취소하고 돌아가기때문에 어느정도 pruning도 된다고 볼 수 있음

 

다만 지금 코드의 경우 문제가 있는데,

 

파란색과 노란색 순서로 진행한 경우. 결과는 동일하지만, 연산을 더 하고있음

 

첨부한 이미지처럼, 0 -> 8 -> 1이나 0 -> 1 -> 8이나 결과적으로 방문한 노드는 동일하지만, 이 '동일 방문'을 거르는 로직이 없기 때문에 이런식으로 다음 노드 후보를 전부 순회하게 되면

 

0 -> 8 -> 1 -> 2 -> ...

0 -> 8 -> 1 -> 4 -> ...

0 -> 8 -> 1 -> 7 -> ...

0 -> 8 -> 1 -> 9 -> ...

0 -> 1 -> 8 -> 2 -> ...

0 -> 1 -> 8 -> 4 -> ...

0 -> 1 -> 8 -> 7 -> ...

0 -> 1 -> 8 -> 9 -> ...

 

이런 방식으로, 계속 같은것이나 다름 없는 후보를 전부 다 순회하게 되며, 불필요한 연산이 늘어나게 되어 낭비가 생김

 

프로그래머스 결과 검증에는 들어있지 않지만 극단적인 테스트 케이스를 살펴보자면

 

노드 최댓값 17, 쫙 퍼진 이진트리, 전부 양

 

info = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

edges = [[0, 1], [0, 2], [1, 3], [1, 4], [3, 7], [3, 8], [4, 9], [4, 10], [2, 5], [2, 6], [5, 11], [5, 12], [6, 13], [6, 14], [14, 15], [14, 16]]

return = 17

 

이런 지극히 극단적인 케이스를 만나게 될 경우 연산은 심히 곤란해지는데, 답은 간단하게 17이지만 늑대가 하나도 없기 때문에 pruning이 전혀 되지 않고, 모든 노드를 전부 후보에 넣으면서 계속 재귀해야하고, 또 순서가 바뀐 노드들 또한 전부 재귀하면서 밑도끝도없이 횟수가 늘어나게 됨. 계산상으로 약 22.4억번의 재귀가 이루어지게 됨

(info를 훑어서 1이 없으면 0 갯수를 그냥 answer로 제출할 수도 있겠지만 그건 넘어간다 치고...)

실제 프로그래머스에서 테스트 케이스로 이걸 넣고 기존 코드를 제출해보면 시간초과로 터지는걸 볼 수 있음

 

그래서 개선점으로 넣은것이 set과 vector<bool>을 통해 '현재까지 방문한 노드 조합'을 이미 방문했는지 검증하는 로직임

 

#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <set>

using namespace std;

//전역 set으로 seen을 추가
set<vector<bool>> seen;

void backtracking(const vector<int>& info, unordered_map<int, vector<int>>& tree, int& answer, int wolf, int sheep, int node, vector<int> candidates, vector<bool> visited){
    int cursheep = info[node] == 0 ? sheep + 1 : sheep;
    int curwolf = info[node] == 1 ? wolf + 1 : wolf;
    if (curwolf >= cursheep) return;
    
    answer = max(cursheep, answer);
    
    //visited 벡터의 현재 노드값 방문을 true로 한 다음, seen에 현재 visited와 동일한 벡터가 있는지 체크
    //동일한 벡터가 있다면 이 노드 조합은 이미 시도한 적이 있기 때문에, 즉시 return함
    //예 : 0 -> 1 -> 8, 0 -> 8 -> 1은 방문 노드가 (0, 1, 8)로 동일하기에 더이상 연산하지 않고 return
    //이후 seen에 현재 노드까지 추가된 visited를 추가
    visited[node] = true;
    if(seen.count(visited)) return;
    seen.insert(visited);
    
    candidates.erase(remove(candidates.begin(), candidates.end(), node), candidates.end());
    for(int i : tree[node])
    {
        candidates.push_back(i);
    }
    
    for(int i : candidates)
    {
        backtracking(info, tree, answer, curwolf, cursheep, i, candidates, visited);
    }
}

int solution(vector<int> info, vector<vector<int>> edges) {
    int answer = 0;
    unordered_map<int, vector<int>> tree;
    for(const vector<int>& v : edges)
    {
        tree[v[0]].push_back(v[1]);
    }
    
    //info의 크기와 동일한 bool 벡터를 선언
    vector<bool> visited(info.size(), false);
    
    backtracking(info, tree, answer, 0, 0, 0, {}, visited);
    
    return answer;
}

 

이 방식을 통하면 이미 방문했던 노드 조합에서는 더이상 재귀가 진행되지 않기 때문에, 메모리는 좀 더 많이 들지만(복사할 vector<bool>이 생겼고 이걸 계속 set에 넣기 때문에) 재귀 횟수가 약 22.4억에서 약 6천번으로 대폭 줄어들게 됨

 

이런식으로 이미 정답처리가 되었다고 안주하지 않고 혹여나 생길 수 있는 문제를 검토해보는것이 개인적으로 많은 도움이 될 것이라 생각하고, 또 재밌기도 하니 종종 해보려고 함