BOJ 백준 1572 : 중앙값

시작하면서

　이번 문제는 세그먼트 트리(Segment Tree) 문제이다.

BOJ 1572 - 중앙값

​

문제 이해

​ 문제에서 K초 이후에는 1초마다 중앙값을 구해야한다. 이걸 단순하게 접근하면 1개의 온도를 추가로 잴 때, K 이전에 쟀던 온도를 1개씩 빼주어야 하고, 이를 정렬한 후에 중앙값을 구해야한다. N이 최대 250,000이므로 O(대략 250,000 번의 시행 * 대략 (250,000 * log(250,000))의 정렬) -> O(250,000 * 250,000 * log(250,000))의 시간복잡도가 나와서 시간 초과가 발생한다.

​ 시간 복잡도를 줄이기 위해서는 다른 방법을 생각해야 한다. 이전의 값들을 그대로 사용하면서 하나씩 업데이트가 되므로 세그먼트 트리를 사용하면 된다. 방법은 다음과 같다.

1. 수의 범위가 0~65536이므로 리프 노드의 범위를 이처럼 잡는다.
2. 인풋을 받으면 배열에 따로 저장한 후, 해당 값을 인덱스로 세그먼트 트리의 리프 노드 값을 +1(카운팅) 해준다.
   • 중앙값을 구하고 나면 배열에 저장해둔 값을 인덱스로 K 이전에 측정했던 온도를 세그먼트 트리에서 지워준다.
3. 중앙값을 구할 때에는 세그먼트 트리에 (N+1)/2을 카운팅한 노드 값으로 찾는다.

​

해결

#include <iostream>
#include <vector>
#include <algorithm>
#include <math.h>

using namespace std;

void UpdateSegmentTree(vector<int>& segmentTree, int node, int index, int start, int end, int diff)
{
    if (index < start || end < index)
        return;

    if (start == end)
    {
        segmentTree[node] += diff;
        return;
    }

    int mid = (start + end) * 0.5f;

    UpdateSegmentTree(segmentTree, node * 2, index, start, mid, diff);
    UpdateSegmentTree(segmentTree, node * 2 + 1, index, mid + 1, end, diff);

    segmentTree[node] = segmentTree[node * 2] + segmentTree[node * 2 + 1];
}

int QuerySegmenTree(vector<int>& segmentTree, int node, int findIndex, int start, int end)
{
    if (start == end)
        return start;

    if (findIndex <= segmentTree[node * 2])
        return QuerySegmenTree(segmentTree, node * 2, findIndex, start, (start + end) * 0.5f);

    return QuerySegmenTree(segmentTree, node * 2 + 1, findIndex - segmentTree[node * 2], (start + end) * 0.5f + 1, end);
}

int main()
{
    // sync off
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    // sync off

    vector<int> segmentTree;
    vector<int> value;
    int N, K, findIndex;
    long long result = 0;
    const int MAX = 65537;

    // segment tree
    int h = (int)ceil(log2(MAX + 1));
    segmentTree.resize(1 << (h + 1));

    cin >> N >> K;

    findIndex = (K + 1) * 0.5f;

    // get value & update segment tree
    value.resize(N + 1);
    for (int i = 0; i < K - 1; i++)
    {
        cin >> value[i];

        UpdateSegmentTree(segmentTree, 1, value[i], 0, MAX, 1);
    }

    for (int i = K - 1; i < N; i++)
    {
        cin >> value[i];

        UpdateSegmentTree(segmentTree, 1, value[i], 0, MAX, 1);
        result += QuerySegmenTree(segmentTree, 1, findIndex, 0, MAX);
        UpdateSegmentTree(segmentTree, 1, value[i - K + 1], 0, MAX, -1);
    }

    cout << result;

    return 0;
}