백준 2957 이진 탐색 트리 C++

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

2957번: 이진 탐색 트리

오랜만에 조금 높은 티어 문제를 풀었습니다. (사실 예전에 틀린 것... 뭔가 최악의 케이스에 걸려 시간초과가 떴었을 듯하다....)

 

문제에서 구하고자 하는 C 값은 사실 이진트리에서 루트부터 몇 개의 노드를 카운트하는 지를 구하는 문제입니다. 단 카운트 수가 누적됨에 주의합니다.

 

문제를 다 풀고 좀 찾아보니 단순 이진 탐색을 이용한 방법 외에 유니온 파인드를 이용해서도 풀 수 있음을 알게 되었습니다. (이 방법도 나중에 시간나면 풀어보는 걸로....)

 

저는 맵을 이용해 단순 이진트리를 통해 문제를 풀었습니다. 맵 자료구조는 내부적으로 키 값을 오름차순으로 정렬시키는 이진 트리 구조를 갖기 때문에 맵 자료구조를 이용하였습니다.

맵의 upper_bound 는 입력된 키보다 큰 값 중 가장 작은 값을 가리키는 iterator를 갖습니다.

예를 들어 맵의 키가 1,2,4 로 주어져 있을 때, 3의 upper_bound는 4를 가리킵니다. 만약 1,2,3,4 로 주어져 있더라도 3의 upper_bound는 4입니다.

 

입력되는 값이 첫 값이면 당연히 루트이고,

이후 입력되는 값이 최대값인 경우, upper_bound가 end를 가리키므로 iterator -1을 통해 이전 iterator의 높이를 구합니다.

최소값인 경우, upper_bound가 begin을 가리키므로 iterator의 높이를 구합니다.

 

이제 가장 애매한 중간의 값이 입력되는 경우, 해당 값이 upper_bound가 가리키는 iter의 왼쪽 자식으로 들어갈 수도 있고, upper_bound의 왼쪽 자식의 오른쪽 자식으로 들어갈 수도 있습니다. 

그러므로 두 케이스 모두 구한 다음, 이진트리이므로 더 깊은 곳으로 들어가게 되므로 높이의 최대값을 구합니다.

 

구해진 높이를 전역으로 설정한 C에 계속 더해가며, 입력된 값을 키로, height+1 값을 value로 하여 맵에 추가합니다.

 

코드는 다음과 같습니다.

 

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>

using namespace std;
// 바이너리 트리 만들면서 체크되는 노드가 몇 개인가를 계산하면됨

int n;
long long C;
map<int,int>m; // 숫자, 뎁스

void init()
{
  cin >> n;
  C = 0;

  // root 만들기
  int num;
  cin >> num;
  m.insert({num, 1});
  cout << 0 << '\n';
  C = 0;
}

void calc(int num)
{
  map<int,int>::iterator iter;
  iter = m.upper_bound(num);

  int height = 0;
  
  // num이 가장 큰 경우
  if(iter == m.end()) {
    C += (--iter)->second;
    cout << C << '\n';
    height = iter->second + 1;
  }
  // num이 가장 작은 경우
  else if(iter == m.begin()) {
    C += iter->second; // upper_bound 이므로 다음 값 가리킴
    cout << C << '\n';
    height = iter->second + 1;
  }
  // 중간인 경우
  else {
    // upper가 가리키는 곳 왼쪽 자식으로 붙을지
    // upper의 왼쪽 자식의 오른쪽 자식으로 붙을지 모름
    // 더 깊은 쪽으로 붙음
    int upperHeight = iter->second;
    iter--;
    int tmpHeight = iter->second;
    height = max(upperHeight, tmpHeight);

    C += height; // 붙을 노드의 height만큼 카운트되는 것과 동일
    cout << C << '\n';
    height++;
  }

  m.insert({num, height}); // +1된 뎁스 추가
}

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

  int num;
  for(int i = 1; i < n; i++) {
    cin >> num;
    calc(num);
  }
}