백준 2243 (사탕 상자) <세그먼트 트리>

문제

2243번: 사탕상자

사탕상자에서 조건에 따라 사탕을 꺼내고 사탕의 맛을 출력하는 문제이다.

사탕의 맛은 1 ~ 100000 이고, 1이 제일 맛있는 맛이다.

입력은 처음에 n이 주어진다. n은 쿼리를 수행할 횟수이다.

이제 n번의 쿼리가 나온다.

a = 1 일경우

a b가 입력으로 들어오고 순위가 b인 사탕으로 꺼내고 해당 사탕의 맛을 출력하는 것이다.

a = 2 일 경우

a b c가 입력으로 들어오고 맛이 b인 사탕 c개를 추가해준다.

 

예를 들어 입력이 아래와 같이 들어왔다면

3
2 1 2
2 2 1
2 3 1

사탕은 순서대로 1123 이 된다.(1순위 = 1, 2순위 = 1, 3순위 = 2, 4순위 = 3)

여기서 우리는 세그먼트 트리를 생각 할 수 있다.

해당 맛의 총 개수를 세그먼트 트리로 저장해 주는 것이다.

위의 입력을 바탕으로 세그먼트 트리를 구현해주면 아래와 같다.

출처 : https://loosie.tistory.com/659

이 문제 같은 경우 "초기의 사탕상자는 빈 상태로 시작한다" 라고 했으므로 별다르게 세그먼트 트리의 노드 값들을 설정해줄 필요가 없다. 

따라서 세그먼트 트리가 만들어질 크기만 계산해서 공간만 할당해주었다.

update(1, 1, 1000000, taste, number);

private static void update(int idx, int s, int e, int taste, int number){
        if(taste < s || taste > e)return;

        tree[idx] += number;

        if(s != e){
            int mid = (s+e)/2;
            update(idx*2, s, mid, taste, number);
            update(idx*2+1, mid+1, e, taste, number);
        }
    }

맛이 1 부터 1000000 까지 임의로 정해주었다. 현재 입력 받은 맛(taste)가 s ~ e 구간에 있으면 모든 세그먼트 트리 노드에 number만큼

더해 주었다.

 

이제 사탕을 꺼내주는 일을 해야 한다. 

만약 순위가 3인 사탕을 꺼내주려면 사탕의 총 개수는 3개 이상인 노드들만 탐색해야 한다.

 

 

당연한것이 나름 포함해서 1번째 부터 3개가 있어야 target을 구할 수 있는 것이다.

왼쪽의 그림을 보면 현재 내가 찾고 있는 순위는 3이기 때문에

먼저 tree[idx*2] 부터 탐색한다.

tree[idx*2] = 3이므로 또 해당 자식 노드들을 살펴볼 수 있다.

왼쪽 자식 노드의 개수는 2이므로 3보다 작기 때문에 바로 오른쪽 노드를 탐색해야 한다.

이 때 target에 tree의 왼쪽 자식을 빼줘야 한다.

 

 

 

 

 

 

 

 

 

이번에는 target = 4를 살펴보자.

 

 

 

처음에 왼쪽 자식 노드는 target보다 작기 때문에 바로 왼쪽 노드로 탐색한다.

이 때 target - tree[idx*2]를 해줘야 한다.

 

 

 

 

 

 

 

 

int candy = query(1,1, 1000000, target);
update(1, 1, 1000000, candy, -1);

private static int query(int idx, int s, int e, int target){
        if(s == e)return s;

        int mid = (s+e)/2;
        if(target <= tree[idx*2])
            return query(2*idx, s, mid, target);
        else return query(2*idx+1, mid+1, e, target-tree[idx*2]);
    }

target이 현재 노드의 자식 노드 보다 작거나 같으면 무조건 왼쪽 자식 노드를 탐색한다.

그게 아니라면 바로 오른쪽 자식 노드를 탐색하고 target - tree[idx*2]를 빼준다.

당연한 것이다. 이미 왼쪽 자식 노드에는 내가 찾는 순위보다 작기 때문에 처음 부터 순위를 차례대로 빼준것과 같기 때문이다.

그 순위를 빼준 값으로 다시 오른쪽 자식노드를 탐색하는 것이다.

 

query를 다 수행하고 target은 하나 꺼냈기 때문에 update를 해줘야 한다. -1 만큼

---

전체 코드이다.

public class Main {
    private static int []tree;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int n = Integer.parseInt(st.nextToken());

        tree = new int[4000004];
        for(int i=0;i<n;i++){
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());

            if(a == 1){
                int target = Integer.parseInt(st.nextToken());
                int candy = query(1,1, 1000000, target);
                System.out.println(candy);
                update(1, 1, 1000000, candy, -1);
            }
            else if(a==2){
                int taste = Integer.parseInt(st.nextToken());
                int number = Integer.parseInt(st.nextToken());
                update(1, 1, 1000000, taste, number);
            }
        }

    }
    private static int query(int idx, int s, int e, int target){
        if(s == e)return s;

        int mid = (s+e)/2;
        if(target <= tree[idx*2])
            return query(2*idx, s, mid, target);
        else return query(2*idx+1, mid+1, e, target-tree[idx*2]);
    }
    private static void update(int idx, int s, int e, int taste, int number){
        if(taste < s || taste > e)return;

        tree[idx] += number;

        if(s != e){
            int mid = (s+e)/2;
            update(idx*2, s, mid, taste, number);
            update(idx*2+1, mid+1, e, taste, number);
        }
    }
}