[BOJ 2243] 사탕 상자, [BOJ 9426] 중앙값 측정

[BOJ 2243] 사탕상자 - P5

문제 설명

1부터 100000까지 라벨이 붙은 사탕이 있다(초기에는 아무 사탕도 없는 상태).
두 가지 쿼리가 들어오는데 첫 번째 쿼리는 k번째로 작은 라벨이 몇 번인지 찾는 것(출력)이고 두 번째 쿼리는 'k번'사탕을 m개 추가하는 것이다.

풀이

[1,1000000]숫자 그대로를 세그먼트 트리의 인덱스로 사용할 것이다. k번째로 작은 라벨이 몇 번인지 찾기 위해서 트리의 i번째 노드는 i번 사탕이 총 몇 개 있는지를 저장하고 있다. 그래서 root 노드에는 결국 전체 사탕의 개수가 담기게 된다. 이제 쿼리에서 루트->리프로 내려갈 때 left child로 내려가면 [1,tree[left]만큼 작은 수가 저장되어 있다는 뜻이다. 따라서 k가 tree[left]보다 크다면 right child로 가야하고 작거나 같다면 left chjld로 가야한다.

소스코드

#include<stdio.h>
#include<algorithm>

using namespace std;
const int SIZE = 1000009;
int tree[SIZE*8];
int startIdx = 1;

void update(int node,int cnt){
    node = node + startIdx - 1;
    tree[node] += cnt;
    int par = node / 2;
    while(par >= 1){
        tree[par] = tree[par*2] + tree[par*2+1];
        par /= 2;
    }
    return;
}

int query(int node,int val){
    if(node >= startIdx){
        return node - startIdx + 1;
    }

    int lCh = node*2,rCh = node*2+1;
    if(tree[lCh] >= val) return query(lCh,val);
    else return query(rCh, val - tree[lCh]);
}

int main(){
    int q;  scanf("%d",&q);
    for(;startIdx < 1000000 ; startIdx *= 2);

    while(q--){
        int cmd;    scanf("%d",&cmd);
        if(cmd==1){
            int x;  scanf("%d",&x);
            int res = query(1,x);
            printf("%d\n",res);
            update(res,-1);
        }

        else{
            int x,y;    scanf("%d %d",&x,&y);
            update(x,y);
        }
    }
    return 0;
}

---

[BOJ 9426] 중앙값 측정 - P5

문제 설명

n개의 수가 주어지고 k개 원소를 가지는 부분수열을 1부터 n-k+1까지 잡았을 때, 부분 배열의 중앙값, 즉 (k+1)/2번째 작은 수를 모두 더해서 출력하는 문제이다.

풀이

[l,r]구간에서 (k+1)/2번 째 작은 수를 log 만에 찾아야 한다. 처음에는 mergesort tree를 써서 제출해보았는데 TLE를 받았다. 아마 n이 25만이라서 터진 것 같다. 그래서 세그먼트 트리로 풀어야 한다.
일단 [l.r]구간에서 k번째 작은 수를 구하는 문제는 비슷한 문제인 k번째 수 문제가 있는데 차이점이라고 한다면, k번째 수 문제는 미리 n개의 수를 모두 입력받아서 쿼리를 받기 전에 mergesort tree로 전처리를 해줘야했지만, 이 문제는 길이가 k인 부분배열만 신경쓰면 되므로 처음부터 굳이 n개의 수를 받아서 전처리 할 필요가 없다. 옆으로 옮길 때마다 삭제되는 값 1개, 추가되는 값 1개로 고정이기 때문에 update함수로 이것만 신경써주고 쿼리는 위에 사탕상자 문제처럼 tree의 노드 값이 현재 갖고 있는 수들의 개수를 가지고 판단하면 된다. k보다 작거나 같으면 left child로, k보다 크면 right child로 가면 된다.

소스코드

#include<stdio.h>
#include<algorithm>
#include<vector>
typedef long long ll;
using namespace std;

const int SIZE = 250009;
int tree[SIZE*8],arr[SIZE];
int startIdx = 1;

int query(int node,int x){
    if(node >= startIdx) return node - startIdx;

    int lCh = node*2,rCh = node*2+1;
    if(tree[lCh] < x) return query(rCh, x-tree[lCh]);
    else return query(lCh,x);
}

void update(int node,int x){
    node = node + startIdx;
    tree[node] += x;
    int par = node / 2;

    while(par >= 1){
        tree[par] = tree[par*2]+tree[par*2+1];
        par /= 2;
    }
    return;
}

int main(){
    int n,k;    scanf("%d %d",&n,&k);

    for(;startIdx < 65536 ; startIdx *= 2);

    for(int i=1 ; i<=n ; i++){
        scanf("%d",&arr[i]);
    }

    ll ans = 0;
    for(int i=1 ; i<=k-1 ; i++)
        update(arr[i],1);

    for(int i=k ; i<=n ; i++){
        update(arr[i], 1);
        ans += query(1,(k+1)/2);
        update(arr[i+1-k], -1);
    }
    printf("%lld\n",ans);

    return 0;
}