[백준 2243번] 파이썬 - 사탕상자

백준 2243 - 사탕상자

시간 제한 2초, 메모리 제한 128MB

# 조건

• 수정이는 어린 동생을 달래기 위해서 사탕을 사용한다.
• 수정이는 평소에 여러 개의 사탕을 사서 사탕상자에 넣어두고, 동생이 말을 잘 들을 때면 그 안에서 사탕을 꺼내서 주곤 한다.
• 각각의 사탕은 그 맛의 좋고 나쁨이 1부터 1,000,000까지의 정수로 구분된다.
• 1이 가장 맛있는 사탕을 의미하며, 1,000,000은 가장 맛없는 사탕을 의미한다.
• 수정이는 동생이 말을 잘 들은 정도에 따라서, 사탕상자 안에 있는 사탕들 중 몇 번째로 맛있는 사탕을 꺼내주곤 한다.
  • 예를 들어 말을 매우 잘 들었을 때에는 사탕상자에서 가장 맛있는 사탕을 꺼내주고, 말을 조금 잘 들었을 때에는 사탕상자에서 여섯 번째로 맛있는 사탕을 꺼내주는 식이다.
• 수정이가 보관하고 있는 사탕은 매우 많기 때문에 매번 사탕상자를 뒤져서 꺼낼 사탕을 골라내는 일은 매우 어렵다.
• 수정이를 도와주는 프로그램을 작성하시오.

입력

• 첫째 줄에 수정이가 사탕상자에 손을 댄 횟수 n(1 ≤ n ≤ 100,000)이 주어진다.
• 다음 n개의 줄에는 두 정수 A, B, 혹은 세 정수 A, B, C가 주어진다.
• A가 1인 경우는 사탕상자에서 사탕을 꺼내는 경우이다.
  • 이때에는 한 정수만 주어지며, B는 꺼낼 사탕의 순위를 의미한다.
  • 이 경우 사탕상자에서 한 개의 사탕이 꺼내지게 된다.
• 또, A가 2인 경우는 사탕을 넣는 경우이다.
  • 이때에는 두 정수가 주어지는데, B는 넣을 사탕의 맛을 나타내는 정수이고 C는 그러한 사탕의 개수이다.
• C가 양수일 경우에는 사탕을 넣는 경우이고, 음수일 경우에는 빼는 경우이다.
• 맨 처음에는 빈 사탕상자에서 시작한다고 가정하며, 사탕의 총 개수는 2,000,000,000을 넘지 않는다.
• 또한 없는 사탕을 꺼내는 경우와 같은 잘못된 입력은 주어지지 않는다.

출력

• A가 1인 모든 입력에 대해서, 꺼낼 사탕의 맛의 번호를 출력한다.

# 접근 방법

• 1 ~ 1,000,000 까지의 맛이 있는 사탕을 저장해주어야 된다.
• 따라서, 맛을 리프 노드로 하여 각 노드를 자신의 담당 구역의 합으로 정의한 세그먼트 트리를 이용해준다.
• 이후 K번째 원소를 찾아야 하므로 아래 로직을 수행해준다.
  • 왼쪽 자식의 값이 K보다 크거나 같다면, 왼쪽 자식으로 들어가서 K번째를 찾는다.
  • 반대로 작다면, 오른쪽 자식으로 들어가서 찾는다.
• 업데이트 또한, 값을 넣을 때는 + 개수만큼, K번째 원소를 찾은 뒤에는 해당 인덱스에 -1을 하면서 update를 해주면 된다.
• 또한, 리프노드에 도착 했을 떄, 해당 노드가 책임지는 사탕의 맛이 무엇인지 필요했기에 따로 index를 저장하는 arr을 만들어 준다.
• 세그먼트 트리에 대한 참고 글
  • https://cheon2308.tistory.com/entry/%EC%9E%90%EB%A3%8C%EA%B5%AC%EC%A1%B0-%EC%84%B8%EA%B7%B8%EB%A8%BC%ED%8A%B8-%ED%8A%B8%EB%A6%ACSegment-Tree

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
from math import ceil, log  

def put_candy(idx, start, end, taste, value):  
    if taste < start or end < taste:  
        return segment_tree[idx]  

    if start == taste == end:  
        segment_tree[idx] += value  
        arr[idx] = taste  
        return segment_tree[idx]  

    mid = (start + end) // 2  
    segment_tree[idx] = put_candy(idx*2, start, mid, taste, value) + put_candy(idx*2+1, mid+1, end, taste, value)  

    return segment_tree[idx]  

def get_candy(idx, start, end, num):  
    if (start == end):  
        segment_tree[idx] -= 1  
        print(arr[idx])  
        return segment_tree[idx]  

    left = segment_tree[idx*2]  
    mid = (start + end) // 2  
    if left >= num:  
        get_candy(idx*2, start, mid, num)  
    else:  
        get_candy(idx*2+1, mid+1, end, num - left)  

    segment_tree[idx] = segment_tree[idx*2] + segment_tree[idx*2 +1]  

    return segment_tree[idx]  



N = int(input())  
H = ceil(log(1000000, 2) + 1)  
tree_size = pow(2, H+1)  

segment_tree = [0] * tree_size  
arr = [0] * tree_size  
for _ in range(N):  
    query = [*map(int, input().split())]  
    if query[0] == 1:  
        get_candy(1, 1, 1000000, query[1])  
    else:  
        put_candy(1, 1, 1000000, query[1], query[2])