[백준 9246번] [세그먼트 트리 / 이분 탐색] 중앙값 측정

Approach

전체 수의 크기가 65536정도밖에 안되는 상황이므로, k개의 숫자들 중 해당 숫자보다 작은 수의 개수를 질의하는 쿼리를 빠르게 응답해주면 된다.

 

잘 생각해보면 세그먼트 트리로 쉽게 구할 수 있다는 것을 알 수 있다.

또한 슬라이딩 윈도우 느낌으로 빠지는 부분은 -1, 들어오는 부분은 +1하는 방식으로 update해주면 된다.

 

추가적으로 위 방식대로 진행하게 되면 0 ~ 특정한 수까지 고려했을 때의 개수가 중앙값 보다 작은 경우와 그렇지 않은 경우가 이분법적으로 갈리게 되므로 이분탐색으로 중앙값을 쉽게 도출할 수 있게 된다.

Code

#include <bits/stdc++.h>
#define fastio cin.tie(0)->sync_with_stdio(0)

using namespace std;
using ll = long long;
using pii = pair<int, int>;
using tiii = tuple<int, int, int>;

int move_x[4] = {-1, 1, 0, 0};
int move_y[4] = {0, 0, -1, 1};

int seg[70000 * 4];

void update(int node, int start, int end, int idx, int value){
    if(idx < start || end < idx) return;
    if(start == end){
        seg[node] += value;
        return;
    }
    seg[node] += value;
    int mid = (start + end) >> 1;
    update(node * 2, start, mid, idx, value);
    update(node * 2 + 1, mid + 1, end, idx, value);
}

ll query(int node, int start, int end, int left, int right){
    if(right < start || end < left) return 0;
    if(left <= start && end <= right) return seg[node];
    int mid = (start + end) >> 1;
    return query(node * 2, start, mid, left, right) + query(node * 2 + 1, mid + 1, end, left, right);
}


int main() {
    fastio;
    memset(seg, 0, sizeof(seg));
    int n, m;
    cin >> n >> m;
    vector<int> v(n + 1);

    for(int i = 1; i <= n; i++){
        cin >> v[i];
        if(i < m) update(1, 0, 70000, v[i], 1);
    }

    ll ret = 0;

    for(int i = m; i <= n; i++){
        update(1, 0, 70000, v[i], 1);
        int start = -1;
        int end = 70000;
        while(start + 1 < end){
            int mid = (start + end) >> 1;
            if(query(1, 0, 70000, 0, mid) < (m + 1) / 2) start = mid;
            else end = mid;
        }
        ret += end;
        update(1, 0, 70000, v[i - m + 1], -1);
    }
    cout << ret << "\n";
    return 0;
}