[ 백준 ] 1572번 - 중앙값 (C++)

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

1572번: 중앙값

 

 

[ 문제풀이 ]

이 문제를 풀기 전에 다음 글을 먼저 읽고 오시는 것을 추천드립니다.

https://rudalsd.tistory.com/51

[ 자료구조 ] 세그먼트 트리 (Segment Tree)

 

1. 초기에 Segment Tree를 초기화해주는 대신 트리를 업데이트해주는 함수만을 통해 K 시간 동안 값들을 업데이트하고, 그중 중앙값을 ans에 더해줍니다.

 

2. K 시간 이후에는 i 초의 값을 세그먼트 트리에 넣어주고, i - K초의 값을 세그먼트 트리에서 빼줍니다.

 

3. 이후 중앙값을 ans에 더해주고, 마지막에 ans를 출력합니다.

 

[ 소스코드 ]

#include<iostream>
#define ll long long
 
using namespace std;
 
int N, K;
int segTree[262222];
int arr[250000];
ll ans;
 
void plusTree(int node, int s, int e, int num)
{
    if (s > num || e < num) return;
    segTree[node]++;
    
    if (s != e) {
        int m = (s + e) / 2;
 
        plusTree(node * 2, s, m, num);
        plusTree(node * 2 + 1, m + 1, e, num);
    }
}
 
void minusTree(int node, int s, int e, int num)
{
    if (s > num || e < num) return;
    segTree[node]--;
 
    if (s != e) {
        int m = (s + e) / 2;
 
        minusTree(node * 2, s, m, num);
        minusTree(node * 2 + 1, m + 1, e, num);
    }
}
 
void getTree(int node, int s, int e, int cnt)
{
    if (s == e) {
        ans += s;
        return;
    }
 
    int m = (s + e) / 2;
 
    if (segTree[node * 2] >= cnt) {
        getTree(node * 2, s, m, cnt);
    }
    else {
        getTree(node * 2 + 1, m + 1, e, cnt - segTree[node * 2]);
    }
}
 
int main()
{
    scanf("%d %d", &N, &K);
    int Max = 0;
    int Min = 66666;
 
    for (int i = 0; i < N; i++) {
        scanf("%d", &arr[i]);
        Max = max(Max, arr[i]);
        Min = min(Min, arr[i]);
    }
 
    for (int i = 0; i < K; i++) {
        plusTree(1, Min, Max, arr[i]);
    }
 
    getTree(1, Min, Max, (K + 1) / 2);
 
    for (int i = K; i < N; i++) {
        plusTree(1, Min, Max, arr[i]);
        minusTree(1, Min, Max, arr[i-K]);
        getTree(1, Min, Max, (K + 1) / 2);
    }
 
    printf("%lld", ans);
}