백준 2243번 사탕상자

문제 링크입니다: https://www.acmicpc.net/problem/2243

세그먼트 트리의 query 부분을 응용해야 풀 수 있는 문제였습니다.

저는 아직 세그먼트 트리 개념이 부족했기 때문에 crocus님의 블로그를 참고하고 풀었습니다.

(https://www.crocus.co.kr/668)

알고리즘은 아래와 같습니다.

1. update 함수는 기존의 세그먼트 트리와 똑같습니다.

2. 사탕의 수는 단말 노드에 있기 때문에 무조건 단말 노드로 향해야합니다.

3. 왼쪽 자식으로 갈지 오른쪽 자식으로 갈지는 해당 노드에 사탕의 누적합이 K개 이상이냐 미만이냐에 따라 나뉩니다.

i) 왼쪽 자식의 누적합이 K개 이상이면 왼쪽 노드로 가면 됩니다.

ii) 왼쪽 자식의 누적합이 K개 미만이면 오른쪽 노드로 가고 (K - 왼쪽 자식의 누적합)번째 사탕을 꺼내면 됩니다.

#include <iostream>

#include <vector>

#include <cmath>

using namespace std;

 

const int MAX = 1000000;

 

int tree_size;

long long result;

 

void update(vector<long long> &tree, int node, int start, int end, int idx, long long diff)

{

        if (!(start <= idx && idx <= end))

                 return;

 

        tree[node] += diff;

 

        if (start != end)

        {

                 int mid = (start + end) / 2;

                 update(tree, node * 2, start, mid, idx, diff);

                 update(tree, node * 2 + 1, mid + 1, end, idx, diff);

        }

}

 

long long query(vector<long long> &tree, int node, int start, int end, int k)

{

        if ((start == end) && !result)

        {

                 cout << start << "\n";

                 return start;

        }

 

        //자신의 왼쪽 자식 노드가 k개 이상 가지고 있다면 왼쪽으로 탐색

        if (result == 0 && (node * 2 <= tree_size && tree[node * 2] >= k))

                 return result = query(tree, node * 2, start, (start + end) / 2, k);

 

        k -= tree[node * 2];

 

        //그 외에는 오른쪽 자식 노드로 탐색

        if (result == 0 && (node * 2 + 1 <= tree_size && tree[node * 2 + 1] >= k))

                 return result = query(tree, node * 2 + 1, (start + end) / 2 + 1, end, k);

}

 

int main(void)

{

        ios_base::sync_with_stdio(0);

        cin.tie(0);

        int N;

        cin >> N;

 

        int height = (int)ceil(log2(MAX));

        tree_size = 1 << (height + 1);

        vector<long long> tree(tree_size);

 

        for (int i = 0; i < N; i++)

        {

                 int num;

                 cin >> num;

 

                 if (num == 1)

                 {

                         int K;

                         cin >> K;

 

                         int idx = query(tree, 1, 0, MAX - 1, K);

                         result = 0;

                         update(tree, 1, 0, MAX - 1, idx, -1);

                 }

                 else if (num == 2)

                 {

                         int rank;

                         long long val;

                         cin >> rank >> val;

 

                         update(tree, 1, 0, MAX - 1, rank, val);

                 }

        }

        return 0;

}

개발환경:Visual Studio 2017

지적, 조언, 질문 환영입니다! 댓글 남겨주세요~