[백준/12899] 데이터 구조 (Python)

Problem : https://www.acmicpc.net/problem/12899

12899번: 데이터 구조

 

Difficulty : Platinum 4

 

Status : Solved(pypy)

 

Time : 00:09:45

 

---

 

문제 설명

 

 

자연수를 저장하는 데이터베이스 S에 대해 다음의 쿼리를 처리합시다.

유형 1 : S에 자연수 X를 추가한다.

유형 2 : S에 포함된 숫자 중 X번째로 작은 수를 응답하고 그 수를 삭제한다.

 

 

입력 및 출력

 

입력

 

첫째 줄에 사전에 있는 쿼리의 수 N 이 주어집니다. (1 ≤ N ≤ 2,000,000)

둘째 줄부터 N개의 줄에 걸쳐 각 쿼리를 나타내는 2개의 정수 T X가 주어집니다.

T가 1이라면 S에 추가할 X가 주어지는 것입니다. (1 ≤ X ≤ 2,000,000)

T가 2라면 X는 S에서 삭제해야 할 몇 번째로 작은 수인지를 나타냅니다. S에 최소 X개의 원소가 있음이 보장됩니다.

 

출력

 

유형 2의 쿼리 개수만큼의 줄에 각 쿼리에 대한 답을 출력합니다.

 

입력 예시

 

5
1 11
1 29
1 89
2 2
2 2

 

출력 예시

 

29
89

 

 

---

 

풀이

 

카운팅 세그먼트트리를 생성하자. X의 최댓값이 2백만이므로, 2백만까지를 수용할 수 있는 세그먼트 트리를 생성한다.

 

• update는 일반적인 세그먼트 트리와 동일하다.
• search의 경우 몇번째 수인지 lower bound를 찾는 경우로 바꾸어 생각할 수 있다.
• 파이썬은 느린 언어. 다시 명심하자. 전부 비재귀 형식으로 구현해 볼 수 있겠다.

 

풀이 코드

import sys
input = sys.stdin.readline
MAX = 2000000
tree = [0]*(4*MAX+4)

def update(target) :
  start, end, idx = 0, MAX, 1
  while start < end :
    mid = (start + end) // 2
    tree[idx] += 1
    if target <= mid :
      idx = idx*2
      end = mid
    else :
      idx = idx*2 + 1
      start = mid + 1
  tree[idx] += 1

def search(target) :
  start, end, idx = 0, MAX, 1
  while start < end :
    mid = (start + end) // 2
    tree[idx] -= 1
    if target <= tree[idx*2] :
      idx = idx*2
      end = mid
    else :
      target -= tree[idx*2]
      idx = idx*2 + 1
      start = mid + 1
  tree[idx] -= 1
  return end

N = int(input())
for _ in range(N) :
  t, x = map(int, input().split())
  if t == 1 :
    update(x)
  else :
    print(search(x))

풀이 완료!