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

문제

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

접근 방법

1. 각 사탕의 맛마다 리프노드에 몇개의 사탕이 있는지 저장.
2. 구간 합을 이용하여 몇번째 맛인지 탐색이 가능하도록 부모 노드들에 저장.
3. 몇번째 맛인지를 이용하여 좌우 자식을 비교하여 이분 탐색을 통해 먹을 사탕 탐색.
4. 먹은 사탕의 맛번호를 이용하여 사탕 상자에서 한개 제거

풀이 코드

import sys
def input(): return sys.stdin.readline().rstrip()

def change(target, diff , idx, start, end):
    if end < target or target < start:
        return 

    tree[idx] += diff
    if start != end:
        change(target, diff , idx*2 , start , (start+end)//2)
        change(target, diff , idx*2+1, (start+end)//2 + 1 , end)

def print_sum(count , idx, start, end):
        
    if start == end: # 리프노드 도달
        return start
    
    else:
        if tree[idx*2] >= count:
            return print_sum(count , idx*2, start, (start+end)//2)
        else:
            return print_sum(count - tree[idx*2] , idx*2 + 1 , (start+end)//2 + 1, end)

    
N = 1_000_000 # 맛 종류
nums = [0] * (N + 1)
tree = [0] * (2**21)

M = int(input())
for _ in range(M):
    order, *content = map(int,input().split())
    if order > 1:
        change(content[0], content[1], 1, 1, N)
        nums[content[0]] += content[1]
        
    else:
        eat = print_sum(content[0] , 1, 1, N)
        print(eat)
        change(eat, -1, 1, 1, N)
        nums[eat] += -1