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

문제링크 : https://www.acmicpc.net/problem/1572

1572번: 중앙값

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define MAX 65536
#define pii pair <int, int>

int N, K;
ll sum_m = 0;
int arr[250001]; 
vector<ll>tree(MAX*4);  //세그먼트 트리

ll sum(int start, int end, int left, int right, int node)  //start ~ end : 노드, left ~ right : 구간
{
    if(left > end || right < start) return 0;  //겹치지 않으므로 종료
    if (left <= start && end <= right) return tree[node];  //교집합이 start, end
    int mid = (start + end) / 2;
    return sum(start, mid, left, right, node*2) + sum(mid+1, end, left, right, node*2+1);
}

void update(int start, int end, int node, int idx, ll diff)
{
    if(idx < start || idx > end) return;  //겹치지 않으므로 종료
    tree[node] = tree[node] + diff; //변한만큼 갱신

    if(start!=end)  //리프노드가 아니라면 자식도 갱신해줌
    {
        int mid = (start+end)/2;
        update(start, mid, node*2, idx, diff);    //왼쪽 탐색
        update(mid+1, end, node*2+1, idx, diff);  //오른쪽 탐색
    }
}

int binary_search()
{
    int l=0, r = MAX;
    int result = -1;
    
    while(l<=r)
    {
        int mid = (l+r)/2;
        
        if(sum(0, MAX, 0, mid, 1)>=(K+1)/2) //K길이의 중앙
        {
            result = mid;
            r = mid -1;
        }
        else
        {
            l = mid + 1;
        }
    }
    return result; 
}

int main(void)
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    cin >> N >> K;
    
    tree.resize(MAX*4);

    for(int i=1; i<=N; i++)
    {
        cin >> arr[i];
        update(0, MAX, 1, arr[i], 1); 

        if(i>=K) //K초부터 이므로
        {
            int m = binary_search(); //중앙값 찾기 
            sum_m += m;
            update(0, MAX, 1, arr[i-K+1], -1); //계산한 부분 트리에서 제거
        }
    }

    cout << sum_m;


    return 0;
}

인생 첫 플래티넘 문제이다.

온전히 내 힘만으로 푼 것이 아니라 구글에 검색 많이하고 따로 분석해가면서 풀 수 있었다.

 

일단 세그먼트 트리와 이분 탐색을 활용한 문제이다.

 

먼저 update를 통해서 1부터 현재 입력값인 arr[i]까지 전부 1을 더해준다.

이를 반복하다가 문제에서 K초부터 표시한다고 했으므로, i가 K보다 같아지는 순간부터는 중앙값을 찾는다.

 

예시를 첫 번째 예제 케이스인

10 3
3
4
5
6
7
8
9
10
11
12

로 보면,

 

i=k인 순간은 3부터이고, arr[3]까지는 계속해서 update로 1 ~ arr[i]까지 1을 더해주는 것을 볼 수 있다.

반복문이 돌떄마다 1씩 더해주는데, arr[3]까지 이므로 총 3번 반복하여 3번 겹치는 부분은 총 3씩 트리에 update 된 것을 알 수 있다.

여기서 arr[2]인 4는 2번 겹치고, arr[3]인 5는 겹치지 않고 1번만 더해진다.

 

이 상태에서 첫 중앙 값을 찾기 시작한다.

우리가 찾는 중앙값은 (K+1)/2인 값을 찾는 것이다.

위의 예제에서는 2다.

 

이분 탐색하는 부분을 보면 여기서 트리의 구간합을 구하고 이 구간합이 2보다 크다면 result에 해당 mid를 담는다.

여기서 left는 가장 작은 값인 0, MAX는 가장 큰 값인 65536이다.

 

현재 트리의 구조는 아래로 갈 수록 겹치는 부분이 줄어든다. (노드가 클수록 값이 작아진다.)

따라서 우리가 찾는 값은 2이기에, 2번만 겹치는 값을 찾는 것이나 다름없다.

처음에는 앞서서 총 3번씩 반복했기에 3인 값만 계속하여 나올 것이다.

 

따로 출력하여 확인해봤는데 2번 겹치는 값인 4는 노드가 16384번이다.

따라서 처음 시작 노드가 1이기에, 굉장히 깊숙히까지 찾는 것을 알 수 있었다.

 

이렇게 찾은 중앙값인 4는 3, 4, 5의 중앙값이다.

그럼 이제 다음으로 찾을 중앙값은 4, 5, 6의 5이다.

 

이에 앞서서 먼저 해주는 작업이 1 ~ arr[i-K+1] 까지 -1을 해줌으로서 더했던 1을 지우는 것이다.

여기서 첫 i가 3이므로 1 ~ arr[1] 즉 맨 처음 값인 1 ~ 3 까지 트리에 1씩 더했던 것을 다시 지우는 작업을 진행한다.

이는 1 ~ 3이 이제는 필요가 없기 떄문이다.

 

1 ~ 3을 제거하고 맨 앞값을 4로 만들어 5가 2번째가 되도록 해준다.

 

이렇게 한번 구했으면 다시 앞에 값을 제거하고 땡겨주는 것을 반복하여 중앙값을 구하고 이때마다 중앙값을 따로 더해준 값을 출력한다.