트리에서 최단 경로를 구하기 위한 또 다른 방법?

 간선의 가중치가 모두 같고, 최단 거리를 구하라 그러면 BFS를 써야 한다. 고 생각하시는 경우가 많습니다. 네. 대부분의 경우에는 맞습니다. 그런데 몇 문제는, 특수한 조건을 줍니다. 이 때에는, BFS를 쓰는 게 답이 아닐 수도 있습니다. 백준 10838번 문제를 봅시다. 일단 솔브드 티어로 골드3. 상당히 만만해 보입니다. 보통 이 경우에는 BFS, DFS 응용 문제들도 많이 분포해 있기 때문에, 최단거리가 나왔다. 그러면 다익스트라나, BFS로 섣불리 판단하기 쉽습니다. 문제는 굉장히 길기 때문에, 따로 설명을 드리지는 않겠습니다. 잘 읽어보셨다면 중요한 조건이 몇 가지 있다는 것을 알 수 있는데요.

 

 트리의 노드 갯수는 10^5개 이하입니다. 여기서 왜 제가 4개의 조건에 밑줄을 쳤을까요? 4개 다 매우 중요한 조건이기 때문입니다. a노드와 b노드의 최단 거리가 1000 이하다라는 것. 1000에 30만을 곱하면 대략 3억이 나오는데요. 이것은 1초에 3억번의 연산을 전제로 계산했을 때 paint와 count를 O(1000)번 정도에 할 수만 있다면, 통과가 된다는 의미입니다.

 

 또한 트리 또한 굉장히 중요한 조건입니다. 이는 제가 10838번 문제를 설명하면서 이야기 해 드리도록 하겠습니다.

 

 

---

 일단 paint와 count 쿼리를 수행하기 위해서는, a와 b의 최단 경로까지 알아내야 합니다. 즉, a에서 출발했을 때, b까지의 최단 경로를 잘 찾아가야 합니다. 이를 BFS로 수행하면 어떻게 될까요? 최단 거리가 1000이니까, 1000개의 노드만 탐색할까요? 잠깐 너비 우선 탐색에 대해서 간단하게 알아봅시다.

 

 

 먼저, 1부터 탐색을 시작할 겁니다. 그러면 1부터 Queue에 들어갈 겁니다. 큐에 들어간 노드를 노란색, 이미 빠진 것을 보라색으로, 방문 아직 안 한 노드를 하늘색으로 표시해 보겠습니다. 먼저 1이 들어갔습니다.

 

 

 그러면 이런 상태가 될 겁니다.

 

 

 다음에 1을 빼 버렸습니다. 1과 인접한 2와 3이 Queue에 들어갈 겁니다.

 

 

 다음에 먼저 들어간 2가 Queue에서 빠지면, 4와 5가 큐에 들어갈 겁니다.

 

 

 3이 빠지면 6이 Queue에 들어간 상태가 됩니다. 1에서 6까지의 최단 거리는 2였습니다. 그런데, 6개의 노드를 모두 visit 해야 했습니다. 그러면, 이런 경우에는 어떨까요? 1과 인접한 노드는 2, ... , 99999입니다. 그리고 99999와 인접한 노드는 10^5입니다.

 

 

 그러면 그래프의 모양이 다음과 같을 겁니다. 1에서 10^5까지 최단 거리는 2입니다. BFS로 탐색하면 어떨까요?

 

 

 일단 1부터 Queue에 들어갑니다. 다음에, 어떻게 되나요? 1이 빠질 겁니다. 다음에 1과 인접한 노드들 모두가 Queue에 들어갈 겁니다. 그러면 2부터 10^5 - 1까지 큐에 들어갈 건데요.

 

 이 과정에서 10^5는 아직 들어오지 않았습니다. 10^5에 대한 정보가 들어오기 위해서는, 10^5 - 1이 큐에서 빠져야 합니다. 

 

 

 물론, 10^5 - 1이 있다면, 자료구조에 넣지 않고 바로 return 해 버리면 시간이 조금은 절약이 될 겁니다. 하지만, 이미, 10^5 - 1개의 노드들은 탐색한 후입니다. 쿼리의 수가 30만개이고, BFS를 돌려서 최단 경로를 찾으려는 경우에는 최악의 경우에 10^5 - 1개의 노드를 방문해야 합니다. 이 둘을 곱하면 300억입니다. 골드 3에 단순하게 BFS를 썼더니 시간 초과. 생각보다 난이도가 높은 게 아닌가 싶습니다.

 

 

---

 문제에서 주어지는 것이 그래프의 일종인 '트리'라고 하였습니다. 트리의 특성을 다시 생각해 봅시다.

 

 이 특성은 이 문제를 푸는 데 상당한 Key가 됩니다. 왜냐하면 역으로 c에서 c의 부모로 간선을 이어버린다면..

 

 

 1에서 10^5로 이어지는 경로는 없지만, 10^5에서 1로 가는 최단 경로는 2라는 것을 매우 빠르게 찾을 수 있기 때문입니다. 이는 단순히, 어떠한 시작점에서 계속 부모를 타고 올라가면서 확인하면 되기 때문이에요. 그런데, 2에서 10^5로 가는 최단 경로는 찾을 수가 없어요. 그런데 쿼리에는 그러한 류의 쿼리가 들어온다는 말입니다. 이건 또 어떻게 해야 할까요?

 

 제가, 방향을 강제로 주었을 뿐. 사실은, 방향성이 없는 Tree입니다. 단지 Root만 정해져 있을 뿐입니다. 그러면, 이렇게 생각해 봅시다. 2에서 계속 부모로 올려봅니다. Root를 만날 때 까지요. 예를 들어 루트가 1이라면 노란색으로 칠해진 노드를 탐색할 겁니다.

 

 

 이렇게 탐색을 한 상태입니다. 다음에 10^5에서부터 부모로 쭉 올라가 볼 건데요.

 

 

1을 최초로 만났습니다. 2와 10^5로부터 가장 가까운 2와 10^5의 공통 조상은 1입니다. 이 때, 1을 2와 10^5의 최소 공통 조상이라고 이야기 합니다. 줄여서 LCA라고 합니다. 이것을 일반화 하면, a에서 a와 b의 LCA로 간 다음에 거기서 b로 가는 경로가 최소 경로라고 말할 수 있어요.

 

 이 LCA는 트리에서 최단 경로를 구할 때 굉장히 중요한 역할을 합니다. 여기서 알고 넘어가셨으면 좋겠습니다.

 

---

 그런데, a로부터 root까지 계속 올라간 다음에, b로부터 root로 올라가는데 이미 a에서 root로 올라가는 과정에서 방문했다면, 그 노드가 a와 b의 최소 공통 조상이라고 체크하는 알고리즘. 효율적일까요? 안타깝지만 root까지 올라가는 것은 그리 효율적이지 않습니다. 예를 들어, a가 10^5이고 b = 10^5 - 1000이고, a와 b의 lca가 10^5 - 1000이라고 해 봅시다.

 

 

 그러면 a에서 root까지 올라가는데 10^5개의 노드를 방문해야 할 거에요. 비효율적입니다. 여기서, a와 b의 최단거리가 1000 이하다라는 조건을 잘 살펴볼 필요가 있는데요. a와 lca까지 거리를 A, b와 lca까지 거리를 B라 했을 때, A + B의 값이 1000 이하입니다. A와 B는 0보다 크거나 같기 때문에, A와 B는 1000 이하가 될 수 밖에 없어요.

 

 

 그렇다면, 많아봐야 1000번 올라가 보면 됩니다. 이를 고려해서 get_lca 함수를 구현하면 위와 같습니다. 남은 것은 노드의 색깔이 아닌, 간선의 색깔인데.. 이것도 크게 어렵지 않습니다.

 

 

 a에서 부모로 올라가는 간선을, a번 간선이라고 넘버링 하면 됩니다. 남은 과제는 count와 paint군요. move는 말 그대로 pa 배열만 업데이트 하면 되는 거고요.

 

---

 color는 [1, 10^9]의 범위로 들어오니, 좌표 압축으로 전처리를 해 주고요. 그렇게 하고 자료구조를 관리해 주면 되는데요. 문제는 색깔의 종류가 최대 10^5개 있을 수 있기 때문에, count 쿼리가 들어올 때 마다 10만개의 크기를 가지는 배열을 초기화 해 주면 시간 초과가 날 수 있습니다. bit로 압축을 하면 이야기가 달라질지도 모르겠지만, 말입니다.

 

 그런데 경로의 최단 거리가 1000 이하라면 나올 수 있는 색깔의 수도 1000 이하입니다. 이는 10만보다 월등히 작은 수입니다. 그러면 count 쿼리가 들어올 때, 방문한 색깔을 visit한 색깔이다. 라고 처리를 해 줍니다. 예를 들어 2, 3, 2를 visit 했다면, 2, 3, 2를 순서대로 자료구조에 넣어줍니다. 다음에, count 쿼리가 끝나면, 자료구조에 있는 색깔들을 0으로 초기화 해 주면 될 거에요.

 

 

 먼저, a에서부터 lca로 올라갑니다. 이 때, color[cur]를 방문하지 않았다면, color[cur]를 방문처리 해 주고, ret 값을 증가시킵니다. 그리고, v에 color[cur]를 넣어줍니다. 이는 방문한 색깔을 의미합니다.

 

 

 

 b에서 lca로 올라갈 때도 마찬가지입니다. 그러면, ret을 구한 다음에, 119 ~ 120번째 줄에서 visit 처리를 초기화 해 줄 수 있습니다. 방문한 색깔 순서가 2, 3, 2였다면, f[2], f[3], f[2]를 모두 0으로 set 해 주면 됩니다.

 

 

 다음에 a에서 b로 이르는 최단 경로의 간선들을 c로 업데이트 하는 것 또한 어렵지 않습니다. 그런데 이건 a가 b의 부모가 아니고, b가 a의 부모가 아닌 경우를 고려를 안 하는 거 아닌가? 라는 생각이 드실 수도 있을 듯 싶습니다. 왜냐하면 이것은, a가 b의 부모일 때만 제대로 동작하는 코드이기 때문입니다. main 함수를 봅시다.

 

 

 get_lca로 a와 b의 lca를 구한 다음에, b에서부터 a로 올라갈 때 까지 탐색하면서 간선의 색깔을 업데이트 하면 됩니다.