BOJ 18240 - 이진 탐색 트리 복원하기

BOJ 18240 - 이진 탐색 트리 복원하기

트리를 직접 잘 구현해서 풀어야 하는 구현 문제이다.

모든 level의 트리를 포인터로 잘 관리한다.

$L$ 레벨의 정점이 들어왔는데 만약 $L-1$ 레벨의 트리 중 더 이상 자식을 넣을게 없다면 정답은 불가능하다.

그렇지 않다면 $L-1$ 레벨 아무데나 일단 붙여서 그런식으로 모두 트리를 구성한다.

이제 만들어진 트리 형태를 역추적하여 현재 루트 노드이고 왼쪽 자식이 $K$ 오른쪽 자식이 $N-1-K$ 개가 있다면 현재 노드의 번호를 $K+1$ 로 배정하는 식으로 하는 것이다.

구현이 좀 더럽긴한데 한 번에 맞긴했다.

struct node {
   node *l = 0, *r = 0;
   int size = 0, idx = 0;
};

void solve() {
   int n;
   cin >> n;
   vi a(n - 1);
   fv(a);

   vector<vector<node *>> by_height(n);
   vi nxt_idx(n);
   node *root = new node();
   root->idx = 0;
   by_height[0].pb(root);
   for (int i = 0; i < n - 1; i++) {
      int h = a[i];
      if (sz(by_height[h - 1]) == 0) {
         cout << -1;
         return;
      }
      int nxt = nxt_idx[h - 1];
      while (nxt_idx[h - 1] < sz(by_height[h - 1]) && by_height[h - 1][nxt_idx[h - 1]]->l != 0
         && by_height[h - 1][nxt_idx[h - 1]]->r != 0) {
         nxt_idx[h - 1]++;
      }
      if (nxt_idx[h - 1] >= sz(by_height[h - 1])) {
         cout << -1;
         return;
      }
      node *new_node = new node();
      new_node->idx = i + 1;
      if (by_height[h - 1][nxt_idx[h - 1]]->l == 0) {
         by_height[h - 1][nxt_idx[h - 1]]->l = new_node;
      } else {
         by_height[h - 1][nxt_idx[h - 1]]->r = new_node;
      }
      by_height[h].pb(new_node);
   }
   vi ans(n);
   function<void(node *)> fn = [&](node *cur) -> void {
      cur->size = 1;
      if (cur->l) {
         fn(cur->l);
         cur->size += cur->l->size;
      }
      if (cur->r) {
         fn(cur->r);
         cur->size += cur->r->size;
      }
   };
   fn(root);
   function<void(node *, int, int)> fn2 = [&](node *cur, int l, int r) -> void {
      assert(cur->size == r - l + 1);
      if (cur->size == 1) {
         ans[cur->idx] = l;
         return;
      }
      ans[cur->idx] = l + cur->l->size;
      fn2(cur->l, l, ans[cur->idx] - 1);
      if (cur->r) fn2(cur->r, ans[cur->idx] + 1, r);
   };
   fn2(root, 0, n - 1);
   for (int i = 0; i < n; i++) cout << ans[i] + 1 << ' ';
}