BOJ 9426번 중앙값 측정

boj.kr/9426

\(1\leq{N}\leq{250000}\)

\(1\leq{K}\leq{5000}\)

\(K\leq{N}\)

\(0\leq{a_i}\leq{65535}\)

이 문제는 N개의 수와 구간의 길이 K가 주어졌을 때 N-K+1 개의 연속된 구간의 중앙값을 모두 더하여

출력하는 문제이다. 중앙값의 정의는 (K+1)/2 번째로 작은 수이다.

naive 하게 생각해보면 O(N)개의 구간에 대해 정렬하여 가운데값을 구하는데 O(NlgN) 이므로 O(N^2lgN)이 든다.

근데 O(N^2)도 시간안에 돌아가기 힘들다.

어떻게할까? 펜윅트리에 수가 현재 길이 K의 구간안에 하나 들어갈 때마다 수 자체를 인덱스로 1업데이트하고

구간에서 사라지면 -1업데이트하는식으로 생각해보자. 그럼일단 0번째인덱스부터 K개의 수에 대해 업데이트해준다.

그럼 현재 펜윅트리에는 각 수가 몇개씩 있는지 query(i)-query(i-1) 로 알 수가 있으며,

현재 구간안에 x이하의 수가 총 몇개가있는지 query(x) 를 통해 알 수가 있다.

그럼 중간값이란 (K+1)/2 번째로 작은 수이기 때문에 K개의 수가 distinct 하다면 

x이하인 수가 (K+1)/2 개 인, x가 바로 중간값일 것이다, 즉 query(x)가 (K+1)/2인 x이다.

query(x) 는 단조증가함수이기 때문에 x값이 커지면 query(x)는 같거나 더 커진다.

그러면 이분탐색으로 query(x)가 (K+1)/2 이 되는 순간의 x를 빠르게 구할 수가 있다는 뜻이 된다.

그럼 만약 K개의 수가 distinct하지 않다면 어떡할까? 약간만 바꿔주면된다.

query(x)가 (K+1)/2 이상으로 바뀌는 순간의 x를 구하면된다. (K+1)/2 로 lower_bound 한다고생각하면되겠다.

(물론 실제로 STL lower_bound를 돌리는건아니고 직접 lo,hi 로 짜는거)

그런데 K구간안에 실제로 들어 있지 않은 수 x에 대해 query(x) 를 날려도 괜찮은지 의문이 들 수도 있다.

우리는 query(x)가 (K+1)/2 이상으로 바뀌는 순간의 첫 x를 찾는 것인데 만약 x가 구간에 존재하지 않는다면

\(query(x-1) = query(x)\) 이기 때문에 상관이 없다. 왜냐하면

query(x-1)이 답이아니라면 query(x)도 같기 때문에 답이아니라고 판단하고 지나가며,

query(x-1)가 답이라면 lower_bound 이기 때문에 상관없다.

그럼 어떤 한 구간에 대한 정보가 펜윅트리에 들어있을 때 O(lg^B) 에(이분탐색X펜윅트리query)

로 중앙값을 구하는건 알았다. (B는 최대 숫자 65535) 그럼 N-K+1개의 구간에 대해 어떻게 할까?

한번 구간을 구했으면 다음 구간은 구간의 맨오른쪽수의 오른쪽에있는 수를 추가 업데이트 시키고

구간의 맨왼쪽수를 삭제 업데이트 시켜주면 결국 다음 구간에 대한 정보가 나온다.(인치웜 알고리즘)

그럼 O(N)에 할 수가 있다. 즉 총 복잡도는 O(Nlg^2B) 가 된다.

머지소트 하면서 정렬된 구간벡터에 쿼리날리듯하면 복잡도가 O(nlgn+n(lg^2n)(lgk)) 이므로 너무느리다.

(트리 구성에 ngln, n개의 구간마다 파라메트릭+query+lower_bound 해서 lgk*lg^2n)

그런데 O(NlgN) 에 할 수 있는 방법이 있다.

바로 K번째 수를 세그트리로 O(nlgn) 에 구할 수 있는 방법인데, 이미 트리에 정보가 업데이트 되어있다면

O(lgn)에 구할 수가 있다. 그냥 정렬해서 K번째인덱스로 접근해도 O(nlgn) 인데 굳이 이런걸 하나 싶을 수도있지만

이렇게 여러 연속된 구간에서 각각을 모두 구해야할 때는 그때 그때 처음부터 벡터를 구성해서 정렬하면 결국

O(n^2lgn) 의 시간이 걸린다. 그리고 여기 뿐만이 아니라 트리에서의 K번째 ~~ 라던지

아무튼 여러곳에 활용가능하다.

그리고 추가적으로 세그트리 뿐만 아니라 이와 비슷한 펜윅 트리에서도 K번째 수를 O(lgn)에 구할 수가 있다.

세그트리와 펜윅트리에서 K번째 수를 찾는 건 여기에서 보도록하자.

아래의 코드는 그냥 펜윅트리 + 이분탐색으로 O(Nlg^2B) 로 짠 소스이다.

#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9+7;
const int INF = 0x3c3c3c3c;
const long long INFL = 0x3c3c3c3c3c3c3c3c;
 
int t[1000003];
int A[1000003];
int N,K;
void update(int pos, int v){
    while(pos < (1<<16)){
        t[pos] += v;
        pos += pos&-pos;
    }
}
int query(int pos){
    int ret = 0;
    while(pos > 0){
        ret += t[pos];
        pos -= pos&-pos;
    }
    return ret;
}
 
int main() {
    scanf("%d%d",&N,&K);
    for(int i = 0 ; i < N; i++) scanf("%d",&A[i]);
    for(int i = 0 ; i < K; i++) update(A[i]+1,1);
    long long sum = 0;
 
    for(int i = K; i <= N; i++){
        int lo = 0, hi = 1<<16;
        while(lo<hi){
            int mid = (lo+hi)/2;
            int s = query(mid);
            if(s >= (K+1)/2) hi = mid;
            else lo = mid+1;
        }
        sum += lo-1;
        if(i==N) break;
        update(A[i]+1,1);
        update(A[i-K]+1,-1);
    }
    printf("%lld",sum);
    return 0;
}
 

이 코드는 K 번째 원소 세그먼트트리의 update 부분만 비재귀(인덱스트리)로 작성한 코드이다. 더 빠르다.

#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9+7;
const int INF = 0x3c3c3c3c;
const long long INFL = 0x3c3c3c3c3c3c3c3c;
 
int t[1000003];
int A[1000003];
int N,K;
int ans;
 
void update(int pos, int v){
    t[pos += (1<<16)] += v;
    for(;pos>1; pos/=2) t[pos/2] = t[pos] + t[pos^1];
}
int search(int nl, int nr, int node, int k){
    if(nl == nr) return nl;
    int nm = (nl + nr)/2;
    if(k <= t[2*node]) return search(nl,nm,2*node,k);
    return search(nm+1,nr,2*node+1,k - t[2*node]);
}
int search(int k){
    return search(0, (1<<16)-1, 1, k);
}
int main() {
    scanf("%d%d",&N,&K);
    for(int i = 0 ; i < N; i++) scanf("%d",&A[i]);
    for(int i = 0 ; i < K; i++) update(A[i],1);
    long long sum = 0;
    for(int i = K; i <= N; i++){
        sum += search((K+1)/2);
        if(i==N) break;
        update(A[i],1);
        update(A[i-K],-1);
    }
    printf("%lld",sum);
    return 0;
}
 

여기서 사실 search 도 비재귀로 짤 수가있다.

#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9+7;
const int INF = 0x3c3c3c3c;
const long long INFL = 0x3c3c3c3c3c3c3c3c;
 
int t[262145];
int A[250003];
int N,K;
int ans;
 
void update(int pos, int v){
    t[pos += (1<<16)] += v;
    for(;pos>1; pos/=2) t[pos/2] = t[pos] + t[pos^1];
}
int search(int k){
    int pos = 1;
    while(pos < (1<<16)){
        if(k <= t[2*pos]) pos*=2;
        else k-=t[2*pos], pos=2*pos+1;
    }
    return pos-(1<<16);
}
 
int main() {
    scanf("%d%d",&N,&K);
    for(int i = 0 ; i < N; i++) scanf("%d",&A[i]);
    for(int i = 0 ; i < K; i++) update(A[i],1);
    long long sum = 0;
    for(int i = K; i <= N; i++){
        sum += search((K+1)/2);
        if(i==N) break;
        update(A[i],1);
        update(A[i-K],-1);
    }
    printf("%lld",sum);
    return 0;
}
 

더빠르다..ㅎㅎ

번외로 BIT도 K번째 원소 구하는게 O(lgn)에 한큐에 된다. 근데 올 비재귀로 짠 세그먼트트리보다 느림 ㅎ

#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9+7;
const int INF = 0x3c3c3c3c;
const long long INFL = 0x3c3c3c3c3c3c3c3c;
 
int t[250003];
int A[250003];
int N,K;
void update(int pos, int v){
    while(pos < (1<<16)){
        t[pos] += v;
        pos += pos&-pos;
    }
}
int search(int k){
    int idx=0;
    for(int i=16;i>=0;i--){
        if(idx+(1<<i)<=(1<<16)-1 && t[idx+(1<<i)] <k){
            k-=t[idx+(1<<i)];
            idx=idx+(1<<i);
        }
    }
    if(!k) return idx;
    return idx+1;
}
int main() {
    scanf("%d%d",&N,&K);
    for(int i = 0 ; i < N; i++) scanf("%d",&A[i]);
    for(int i = 0 ; i < K; i++) update(A[i]+1,1);
    long long sum = 0;
 
    for(int i = K; i <= N; i++){
        sum += search((K+1)/2)-1;
        if(i==N) break;
        update(A[i]+1,1);
        update(A[i-K]+1,-1);
    }
    printf("%lld",sum);
    return 0;
}