Tip - 세그먼트 트리에서 K번 째 수 찾기

존재하는 값의 범위가 $[L, R]$ 이고 어떤 값이 존재할 때마다 트리에 $1$을 더한다고 할 때,

세그먼트 트리에서 $K$ 번째로 작은 수를 찾는다고 해보자.

이는 이분 탐색을 쓰면 $O(\log^2 N)$ 에 해결된다.

특정 수 $M$에 대해 $[0,M]$ 까지 쿼리를 날려보고 $K$개 이상이라면 정답을 $M$으로 설정해주고 이분탐색을 좀 더 작은 $M’$ 에 대해서 돌리거나 반대라면 좀 더 큰 $M’$ 에 대해서 돌릴 수 있기 때문이다.

하지만 세그먼트 트리의 구조를 잘 활용하면 이 쿼리를 ${\color{salmon} O(\log N) }$ 에 돌릴 수 있다.

쿼리 함수를 만들고, 계속해서 왼쪽 자식과 오른쪽 자식이 갖고있는 수 개수를 비교하면서 추적하는 것이다.

주의할 점은 다음과 같다.

• 세그먼트 트리의 특정한 구간 $[l, r]$ 에서 $K$ 번째로 큰 수를 찾는것이 아니다. 그 문제는 Merge Sort Tree나 Persistent Segment Tree로 풀어야 한다.
• 세그먼트 트리를 배열의 구간 대해서 다루는 상태가 아닌 특정 수가 존재하는 개수에 대해서 구성을 했을 때 사용할 수 있다.

현재 $K$ 번째 수를 찾아야 한다고 할 때, 왼쪽 자식의 개수가 $L$ 이고 오른쪽 자식의 개수가 $R$ 이라면,

• $L \ge K$ 라면 그 다음 재귀함수는 왼쪽 자식으로
• 그렇지 않다면 $K \coloneqq K-L$ 을 해준다음 오른쪽 자식으로 이동해준다.
  • 왼쪽 자식만큼 skip 했기 때문이다.

구현은 다음과 같다.

struct Seg {  
   ...  
   int kth(int n, int nl, int nr, int k) {  
      if (nl == nr) return nl;  
      int left = tree[n * 2], right = tree[n * 2 + 1];  
      int m = nl + (nr - nl) / 2;  
      if(left >= k) return kth(n * 2, nl, m, k);  
      else return kth(n * 2 + 1, m + 1, nr, k - left);  
   }  
};