[Python 3] BOJ 2243 : 사탕상자

오늘 풀어볼 문제는 사탕상자 이다.

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

2243번: 사탕상자

문제를 보면, 1부터 100만까지의 정수로 맛의 정도를 구분할 수 있는데, 1이 제일 맛있고, 100만이 제일 맛이 없다. 여기서 수정이는 동생에게 'n'번째로 맛있는 사탕을 꺼내줄 수 있다.

 

사탕상자에서 사탕을 꺼낼 수도 있고, 사탕 상자에 사탕을 넣을 수도 있다. 사탕을 넣는 경우는 n개의 입력 중 1로 시작하는 는 input이 들어온 경우만 가능하다.

 

1부터 100만까지의 사탕을 나이브하게 저장하자면, 최대 100만개의 숫자를 보면서 몇 번째 사탕인지를 찾아야할 수도 있다. 즉, 이를 10만번 해도, 100만 * 10만 개의 연산을 해야할 수도 있을 것이고, 이러면 시간 제한 내에 계산을 할 수 없을 것이다.

 

그렇기에, Segment Tree를 응용하여 문제를 해결하였다.

 

 

n=int(input())

data=[0]*(10**6+1)
tree=[0]*(2**22)

def changeval(target, difference, index, start, end):
    if end<target or target<start:
        return
    tree[index]+=difference
    if start!=end:
        changeval(target,difference,index*2,start,(start+end)//2)
        changeval(target,difference, index*2+1, (start+end)//2+1,end)

def printval(count,index,start,end):
    if start==end:
        return end
    else:
        if tree[index*2]>=count:
            return printval(count,index*2,start,(start+end)//2)
        else:
            return printval(count-tree[index*2],index*2+1,(start+end)//2+1,end)

for _ in range(n):
    a,*b=map(int,input().split())
    
    if a==1:
        b=b[0]
        answer=printval(b,1,1,10**6)
        print(answer)
        changeval(answer,-1,1,1,10**6)
        data[answer]-=1

    else:
        b,c=b
        changeval(b,c,1,1,10**6)
        data[b]+=c

 

일반적인 캔디의 개수를 세는 data와 세그먼트 트리인 tree를 선언하였다. 1은 printval과 changeval을, 2는 changeval을 통해 for문을 실행하였다. 

changeval에서는 target, 즉 내가 원하는 target이 start index와 end index 사이에 있다면, tree에 대한 index를 계속 변경시킨다. 

printval은 treeval의 index를 발견시킨다. 

printval에서는 몇번째 index가 b번째 index인지를 출력해준다.