[segment tree] 2243번 사탕상자

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

자료구조

• Segment Tree

사용 언어

• Python

시간 복잡도

• $O(10^5 log(10^6))$

풀이

포인트

• 트리의 i번째 노드 값이 1~i 사탕 맛의 개수의 합을 의미하도록 트리 설정
• 왼쪽 자식 노드의 값과 오른쪽 자식 노드의 값을 k와 비교하면서 왼쪽 자식 노드와 오른쪽 자식 노드 중 하나를 선택함

Segment Tree

이 문제에서 segment tree 자료구조를 사용해야하는 이유에 대해서 생각해보자. 굳이 자료구조를 사용해서 어렵게 풀지 않아도 될 상황이면, 머리 아프게 자료구조를 생각하지 않아도 되기 때문이다.
자료구조를 사용하지 않고 평범한 for문을 돌렸을 때, 시간 복잡도가 어떻게 나오는지 생각해보자. 수정이가 사탕상자에 손을 댄 횟수는 최대 100,000회이고, 손을 댈 때마다 사탕을 꺼내거나 사탕을 넣는다. 길이가 1,000,000(최대 사탕 맛)인 리스트 box을 생각해보자. box[i]는 i 사탕 맛이 현재 몇개가 있는지를 의미한다. 그러면, k 사탕 맛을 넣을 때는 리스트의 k번째 원소의 값을 추가하면 되므로 $O(1)$의 시간이 소요된다. 하지만 b번째로 맛있는 사탕을 꺼내기 위해서는 box 리스트의 처음부터 원소를 방문해야한다. 즉, 최악의 경우 $O(10^6)$의 시간이 소요된다. 종합하면, 전체 시간 복잡도는 $O(10^{10})$이므로 TIL이 발생한다.

즉, b번째로 맛있는 사탕을 꺼내는 행위를 O(logn)의 시간으로 처리해야하므로 트리 자료구조를 사용해야하는 motivation을 얻을 수 있다.

이제, 본격적으로 segment tree을 짜보자. 트리의 i번째 노드 값이 1~i 사탕 맛의 개수를 의미하도록 트리를 구성한다. 이렇게 구성하는 이유는, 보통의 구간합을 구하는 segment tree 문제와는 다르게 query 함수를 구성할 것이기 때문이다.

BOJ 2042번과 같은 문제에서 query 함수는 아래와 같다.

def query(node, start, end, left, right):
    if left <= start and end <= right:
        return tree[node]
    if end < left or right < start:
        return 0
    return query(node*2, start, (start+end)//2, left, right) + query(node*2 + 1, (start+end)//2 + 1, end, left, right)

즉, 현재 노드에서 왼쪽 자식 노드와 오른쪽 자식 노드를 모두 방문하고 그 값을 더한다. 구간 합을 구해야하기 때문에 두 자식 노드를 모두 방문하는 것이다. 하지만 사탕상자 문제에서는 두 자식 노드 중 하나를 아래와 같이 선택적으로 방문한다.

def query(node, start, end, k):
    if start == end:
        return start
    if tree[node*2] >= k:
        return query(node*2, start, (start+end)//2, k)
    else:
        return query(node*2 + 1, (start+end)//2 + 1, end, k-tree[node*2])

여기서 tree[node*2] >= k가 도대체 어떤 의미인지 아래 그림으로 이해해보자.


5번째로 맛있는 사탕을 꺼낸다고 생각해보자. 트리 노드 값에 따라서 아래와 같이 탐색을 해보자.

• 0~8번째로 맛있는 사탕 개수의 합은 15이다.
• 0 ~ 4번째로 맛잇는 사탕 개수의 합은 8이다. 우리는 5번째로 맛있는 사탕을 꺼내야하므로 0 ~ 4번째로 맛있는 사탕 중 하나가 5번째로 맛있는 사탕일 것이다. 왜냐하면 0 ~ 4번째로 맛있는 사탕 개수의 합인 8이 5보다 크기 때문이다.
• 5번째로 맛있는 사탕이 왼쪽 자식 노드에 있다는 것이 확실하기 때문에 오른쪽 자식 노드는 방문하지 않는다. (if-else 구문)
• 바로 이런 상황에서, if tree[node*2] >= k에 걸리게 되어서 오른쪽 자식 노드는 방문하지 않고, 왼쪽 자식 노드만 방문한다. (return query(node*2, start, (start+end)//2))

다음으로, tree[node*2] >= k에 걸리지 않는 경우를 아래 그림으로 이해해보자.


이번에는 12번째로 맛있는 사탕을 꺼내는 상황이다. 트리 노드 값에 따라서 아래와 같이 탐색을 해보자.

• 0~8번째로 맛있는 사탕 개수의 합은 15이다.
• 0 ~ 4번째로 맛있는 사탕 중에는 12번째로 맛있는 사탕이 존재하지 않는다. 0 ~ 4번째로 맛있는 사탕 개수의 합은 8개이므로 최대 8번째로 맛있는 사탕까지만 존재하기 때문이다. 따라서 왼쪽 자식 노드는 방문하지 않는다. (if문을 방문하지 않고 else문을 방문)
• 이제 우리는 5 ~ 8번째로 맛있는 사탕 중에 12번째로 맛있는 사탕이 존재한다는 것을 안다. 근데 잘 생각해보면 12번째로 맛있는 사탕은 왼쪽 자식 노드의 0 ~ 4번째로 맛있는 사탕 개수의 합을 모두 포함한다.
• 그런데 오른쪽 자식 노드는 0번째가 아니라 5번째부터 시작하는, 5 ~ 8번째로 맛있는 사탕 개수의 합이다. 따라서 시작 부분을 맞춰줄 필요가 있다.
• 이러한 이유로 오른쪽 자식 노드로 갈 때는 k에서 tree[node*2] 값을 빼주는 것이다.

오른쪽 자식 노드로 갔으면, 이후에는 재귀적으로 동일한 과정을 반복한다.



시간 복잡도에 대해서 생각해보자. 총 100,000번의 쿼리가 수행되고 각 쿼리마다 세그먼트 트리를 방문하므로 $log(10^6)$의 시간이 소요된다. 따라서 총
$O(10^5 log(10^6))$의 시간이 소요되어 파이썬으로 시간 내에 문제를 해결할 수 있다.

[Source Code]

import sys
input = sys.stdin.readline

def subsum(node, start, end, k):
    if start == end:
        # print(start+1)
        return start
    if tree[node*2] >= k:
        return subsum(node*2, start, (start+end)//2, k)
    else:
        return subsum(node*2 + 1, (start+end)//2 + 1, end, k-tree[node*2])

def update(node, start, end, index, diff):
    if not (start <= index <= end):
        return
    tree[node] += diff
    if start != end:
        update(node*2, start, (start+end)//2, index, diff)
        update(node*2 + 1, (start+end)//2 + 1, end, index, diff)

MAX = 1000000
n = int(input())
tree = [0]*(MAX*4)
tree_size = len(tree)
l = [0]*(MAX)

for _ in range(n):
    cur = list(map(int,input().split()))
    ret = 0
    if cur[0] == 1:
        _, b = cur
        idx = subsum(1,0,MAX-1,b)
        print(idx+1)
        l[idx] -= 1
        update(1,0,MAX-1,idx,-1)
    if cur[0] == 2:
        _, b, c = cur
        update(1,0,MAX-1,b-1,c)