문제 설명
이진 검색 트리(BST)와 하나의 목표값(target)이 주어졌을 때, 트리에 있는 값 중 target에 가장 가까운 k개의 값을 찾아야 합니다. 이때 target은 부동소수점 실수라는 점에 유의해야 하며, k는 항상 유효한 값(k ≤ 전체 노드 수)이라고 가정할 수 있습니다.
예를 들어 아래와 같은 트리가 주어지고,
target = 3.714286, k = 2라면 출력은 [4, 3]이 됩니다.
알고리즘 접근 방식
이 문제는 두 개의 스택을 활용하면 효율적으로 해결할 수 있습니다. 하나는 target보다 작은 값들을 관리하는 smaller 스택, 다른 하나는 target보다 크거나 같은 값들을 관리하는 larger 스택입니다. 각 단계마다 두 스택의 최상단(top) 값 중 target에 더 가까운 것을 결과에 추가하고, 선택된 노드의 서브트리를 계속 탐색하여 다음 후보를 준비합니다.
구체적인 풀이 단계는 다음과 같습니다.
pushSmaller() 함수를 정의합니다. 이 함수는 노드(node), 스택(st), 목표값(target)을 매개변수로 받습니다.
node가 NULL이 아닌 동안 다음을 반복합니다.
node의 값이 target보다 작으면:
node를 st에 삽입(push)
node를 node의 오른쪽 자식으로 갱신
그렇지 않으면:
node를 node의 왼쪽 자식으로 갱신
pushLarger() 함수를 정의합니다. 마찬가지로 노드(node), 스택(st), 목표값(target)을 받습니다.
node가 NULL이 아닌 동안 다음을 반복합니다.
node의 값이 target보다 크거나 같으면:
node를 st에 삽입(push)
node를 node의 왼쪽 자식으로 갱신
그렇지 않으면:
node를 node의 오른쪽 자식으로 갱신
메인 메서드에서는 다음을 수행합니다.
결과를 담을 배열 ret을 선언합니다.
스택 smaller와 larger를 각각 선언합니다.
pushLarger(root, larger, target)을 호출합니다.
pushSmaller(root, smaller, target)을 호출합니다.
k번 반복하면서(매 회차마다 k를 1씩 감소) 다음을 수행합니다.
smaller가 비어 있지 않고, (larger가 비어 있거나 |target − smaller.top()->val| < |target − larger.top()->val|)인 경우:
curr = smaller.top()으로 설정한 뒤 smaller에서 제거(pop)
curr의 값을 ret의 끝에 추가
pushSmaller(curr->left, smaller, target)을 호출
그렇지 않으면:
curr = larger.top()으로 설정한 뒤 larger에서 제거(pop)
curr의 값을 ret의 끝에 추가
pushLarger(curr->right, larger, target)을 호출
ret을 반환합니다.
작동 원리
pushSmaller는 target 미만의 후보들을, pushLarger는 target 이상의 후보들을 중위 순회(in-order) 순서대로 스택에 쌓습니다. 덕분에 smaller 스택의 top은 항상 target보다 작은 값 중 가장 큰 값(전임자, predecessor)이 되고, larger 스택의 top은 target 이상인 값 중 가장 작은 값(후속자, successor)이 됩니다. 따라서 매 단계마다 이 두 값만 비교하면 target에 가장 가까운 값을 빠르게 골라낼 수 있습니다.
C++ 구현 예제
아래 구현을 통해 더 잘 이해해 보겠습니다.
#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
cout << "[";
for(int i = 0; i<v.size(); i++){
cout << v[i] << ", ";
}
cout << "]"<<endl;
}
class TreeNode{
public:
int val;
TreeNode *left, *right;
TreeNode(int data){
val = data;
left = NULL;
right = NULL;
}
};
void insert(TreeNode **root, int val){
queue<TreeNode*> q;
q.push(*root);
while(q.size()){
TreeNode *temp = q.front();
q.pop();
if(!temp->left){
if(val != NULL)
temp->left = new TreeNode(val);
else
temp->left = new TreeNode(0);
return;
}
else{
q.push(temp->left);
}
if(!temp->right){
if(val != NULL)
temp->right = new TreeNode(val);
else
temp->right = new TreeNode(0);
return;
}
else{
q.push(temp->right);
}
}
}
TreeNode *make_tree(vector<int> v){
TreeNode *root = new TreeNode(v[0]);
for(int i = 1; i<v.size(); i++){
insert(&root, v[i]);
}
return root;
}
class Solution {
public:
vector<int> closestKValues(TreeNode* root, double target, int k) {
vector<int> ret;
stack<TreeNode*> smaller;
stack<TreeNode*> larger;
pushLarger(root, larger, target);
pushSmaller(root, smaller, target);
while (k--) {
if (!smaller.empty() && (larger.empty() || (abs(target - smaller.top()->val) < abs(target - larger.top()->val)))) {
TreeNode* curr = smaller.top();
smaller.pop();
ret.push_back(curr->val);
pushSmaller(curr->left, smaller, target);
}
else {
TreeNode* curr = larger.top();
larger.pop();
ret.push_back(curr->val);
pushLarger(curr->right, larger, target);
}
}
return ret;
}
void pushSmaller(TreeNode* node, stack <TreeNode*>& st, double target){
while (node) {
if (node->val < target) {
st.push(node);
node = node->right;
}
else {
node = node->left;
}
}
}
void pushLarger(TreeNode* node, stack <TreeNode*>& st, double target){
while (node) {
if (node->val >= target) {
st.push(node);
node = node->left;
}
else
node = node->right;
}
}
};
main(){
Solution ob;
vector<int> v = {4,2,5,1,3};
TreeNode *root = make_tree(v);
print_vector(ob.closestKValues(root, 3.7142, 2));
}
실행 결과
입력
{4,2,5,1,3}, 3.7142, 2
출력
[4, 3]
복잡도 분석
균형 잡힌 BST에서는 초기 스택 구성에 O(log n)이 소요되고, 이후 k개의 값을 하나씩 추출할 때마다 다음 전임자/후속자를 찾는 데 O(log n)이 걸리므로 전체 시간 복잡도는 약 O(k log n)입니다. 트리가 한쪽으로 치우친 최악의 경우에는 O(n)까지 증가할 수 있습니다. 공간 복잡도는 스택에 저장되는 노드 수에 비례하여 O(n)입니다.