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

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

2243번: 사탕상자

 

학교 동아리 스터디에서 세그트리 연습문제로 나와 풀게 되었다.

스터디 수업때 다뤘던 연습문제와 세그트리를 적용하는 결이 달라 한참 고민을 하고 풀었다.

 

보통 '세그먼트 트리' 하면 유명한 연습문제가 '구간 합 구하기' 이다.

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

2042번: 구간 합 구하기

이 문제의 경우 주어진 수열에 대해 '특정 구간' 을 잡고 그 구간의 합을 구하면 된다.

따라서 주어진 수열에 따라 미리 세그트리를 설정해두고, 숫자가 바뀌면 업데이트, 구간 조회를 요청하면 요청 결과를 세그 트리 내 각 구간 별 데이터를 잘 조합해 돌려주면 되었다.

 

그러나 사탕상자 문제는 구간 합 구하기 문제와 다른 점이 있었다.

우선 구간 합 구하기 문제의 경우 '구하는 것' 이 문제 제목에도 명시되어 있듯이 '구간 합' 이므로 세그 트리에 '구간 합' 을 저장하면 되었다.

리프 노드에는 각 수열의 구성 숫자가 들어가는데, 이 역시 재귀적으로는 크기가 1인 '구간의 합' 과 같으니 세그 트리 구성시에 이해하는 난이도가 높지 않았다.

하지만 사탕상자 문제는 생각할 요소가 많아서 무엇을 저장해야 하는지 결정하기가 힘들었다.

 

우선 '우선 순위가 높은' 사탕을 돌려주어야 하니, 각 사탕의 우선순위를 고려하여 저장해야 할 것 같은 느낌이 들었다.

세그먼트 트리가 BST와 정말 헷갈리는 게, 세그트리는 각 노드 위치의 의미가 '구간' 이고, BST는 각 노드의 위치 의미가 '순서' 라서 '우선 순위' 를 생각 하면서 아이디어를 생각하니 자연스럽게 세그먼트 트리가 아니라 BST로 구현을 하려고 하게 되었다.

 

또한 각 우선 순위 (맛) 별로 사탕을 추가하고 빼는 과정이 있다는 점이 더 혼란스럽게 했다.

사탕의 갯수가 증가하고 감소하는 입력을 보면서 마치 힙에다 데이터를 추가하고 빼는 과정을 연상하게 되니 머릿속에서 세그먼트 트리, BST, 힙이 왔다갔다 하면서 더 혼란스러웠다.

---

그러다가 세그먼트 트리의 특징인 '구간 쿼리' 에 좀 더 집중해서 어떻게 해야 '구간' 을 설정할 수 있을지 고민해보기 시작했다.

그러자 한가지 아이디어가 떠올랐는데, 바로 '사탕 맛' 의 구간을 잡는 것이다.

 

루트 노드를 1번째 맛부터 100만번쨰 맛까지 사탕의 개수의 총 합으로 생각하자.

그리고 루트 노드의 왼쪽 노드를 1번째 맛부터 50만번째 맛까지 사탕의 개수의 총합으로

루트 노드의 오른쪽 노드를 50만 + 1번째 맛부터 100만번째 맛까지 사탕의 개수의 총 합으로 생각하자.

이런 식으로 반씩 구간을 나눠 사탕 갯수의 총 합을 쭉 저장하는 구조로 생각하고 세그트리를 짜면 된다.

 

이 문제는 '구간 합 구하기' 문제와 달리 전체 수열이 주어지지 않고, 나중에 사탕이 추가되는 방식이라 헷갈린다.

이는 간단하게 초기에 각 맛 별 사탕의 개수가 0개인 것으로 생각하면 된다.

그리고 사탕이 추가되면, 해당 맛의 사탕 갯수를 그 추가된/빼는 정도만큼 바꾸도록 업데이트 한 뒤, 이를 전체 구간합 노드에 반영해주면 된다.

---

쿼리를 할 때는, 마치 이분탐색을 하듯이 생각하면 된다.

만약 1부터 50만번째 사탕의 개수가 25개, 50만+1부터 100만까지 사탕의 개수가 25개라고 해보자.

(1부터 100만까지 사탕의 개수는 50개이다.)

 

이 상황에서 27번째 사탕은 어떻게 찾을 수 있을까?

27은 50이하 이므로 27순위의 사탕은 1부터 100만번째 맛 범위 안에 안에 있다.

 

이제 왼쪽 서브트리를 보자.

왼쪽 서브트리가 나타내는 구간인 1부터 50만까지 범위에는 사탕이 총 25개 있다.

즉, 1등부터 25등까지가 왼쪽 서브트리에 있는 것이다.

 

27번째는 27등이므로 절대 왼쪽에 있을 수 없다.

따라서 오른쪽을 탐색한다.

 

오른쪽 서브트리에는 (그 부모 노드의 등수 기준) 26등부터 50등까지가 존재한다.

하지만 나중에 오른쪽 서브트리를 탐색할 때는, 그 오른쪽 서브트리 기준으로 1등부터 25등 사이를 탐색해야 한다.

따라서 구하는 등수 27에서 기존 왼쪽 서브트리에 들어있던 사탕 개수인 25를 빼준다.

그러면 자연스럽게 오른쪽 서브트리에서 27 - 25 = 2등인 사탕을 구하게 된다.

 

이를 재귀적으로 수행하면 우리가 원하는 맛의 사탕을 찾을 수 있다.

 

원하는 맛의 사탕을 찾은 뒤, 그 사탕을 '빼야' 하므로, 다시 해당 맛의 사탕 개수를 1개 빼주도록 업데이트 하면 사탕을 찾는 문제 요구도 해결된다.

---

이를 소스코드로 구현하면 아래와 같다.

import sys
input = sys.stdin.readline


def update(node, start, end, taste, diff):
    if taste < start or end < taste:
        return seg_tree[node]

    if start == end:
        # start == end == taste:
        seg_tree[node] += diff
        return seg_tree[node]

    mid = (start + end) // 2
    left = update(node*2, start, mid, taste, diff)
    right = update(node*2+1, mid+1, end, taste, diff)

    seg_tree[node] = left + right
    return seg_tree[node]


def query(node, start, end, rank):
    if start == end:
        return start

    mid = (start + end) // 2
    # start ~ mid 범위에 찾는 등수가 있다면 왼쪽만 탐사 / 아니라면 오른쪽을 탐사
    if rank <= seg_tree[node*2]:
        return query(node*2, start, mid, rank)
    return query(node*2+1, mid+1, end, rank - seg_tree[node*2])

# (0하고) 1부터 100만까지 100만(+1)개의 리프노드가 존재한다.
seg_tree = [ 0 for _ in range(4000001)]

n = int(input())
for _ in range(n):
    a, *b = map(int, input().split())
    if a == 1:
        taste = query(1, 1, 1000000, b[0])
        print(taste)
        update(1, 1, 1000000, taste, -1)
    else:
        update(1, 1, 1000000, b[0], b[1])

 

이 문제는 중복된 수가 존재하는 정렬된 수열에서 특정 숫자들을 추가하거나 뺐을 때 정렬을 유지하면서 원하는 시점에서 특정 위치의 숫자를 판별하는 문제와 같다.