[백준] 2243 - 사탕상자 (Python)

사탕에는 $1$에서 $1,000,000$까지 맛이 있다. 그렇다면 길이가 $1,000,000$인 길이를 갖는 리스트로 구현이 가능하다. 인덱스는 사탕의 맛이 되고, 값은 사탕의 개수가 될 것이다.

하지만 이렇게만 구현하면 순위로 사탕을 빼는 것이 매우 비효율적인 연산이 된다. 매번 처음부터 사탕의 개수를 세어야 하기 때문이다. 결국 구간별로 사탕의 개수를 빠르게 확인 가능한 세그먼트 트리를 사용해야 한다.

루트에는 모든 사탕 개수의 합이 들어가고 리프에는 각각의 맛 하나의 사탕의 개수가 들어가는 세그먼트 트리를 생각하자.

import sys
input = sys.stdin.readline

n = int(input())
tree = [0]*(4*10**6) # 세그먼트 트리

def putCandy(left, right, node, b, c):
# b => 사탕 맛, c => 사탕 개수
    tree[node] += c
    if left == right:
        return
    mid = (left+right) // 2
    if b <= mid:
    # 맛이 왼쪽 구간에 포함된 경우
        putCandy(left, mid, node*2, b, c)
    else:
    # 맛이 오른쪽 구간에 포함된 경우
        putCandy(mid+1, right, node*2+1, b, c)
    
def findCandy(left, right, node, b):
# b => 순위
    tree[node] -= 1 
    if left == right:
    # 리프노드
        return left # 빼내는 사탕의 맛
    mid = (left+right) // 2 
    if tree[node*2] >= b:
    # 원하는 사탕이 왼쪽 구간에 있음
        return findCandy(left, mid, node*2, b)
    else:
    # 왼쪽 구간의 사탕 개수만큼 순위에서 제거
        return findCandy(mid+1, right, node*2+1, b-tree[node*2])
    
for _ in range(n):
    line = list(map(int, input().rstrip().split()))
    if len(line) == 2:
        b = line[1]
        print(findCandy(1, 10**6, 1, b))
    elif len(line) == 3:
        _, b, c = line
        putCandy(1, 10**6, 1, b, c)